Baoyu Danger Gemini Web

{
  "lockfileVersion": 1,
  "configVersion": 1,
  "workspaces": {
    "": {
      "name": "baoyu-danger-gemini-web-scripts",
      "dependencies": {
        "baoyu-chrome-cdp": "^0.1.1",
      },
    },
  },
  "packages": {
    "baoyu-chrome-cdp": ["baoyu-chrome-cdp@0.1.1", "", {}, "sha512-OR3PQ7NzJDykCXl20TnkZuwvNQ0hsVC2czje93P72xQaA3vKOyPN/Q1CwEgKuYzP7Rka4Fdh4HvURj6AoNR7Tg=="],
  }
}

import assert from "node:assert/strict";
import test from "node:test";

import { collect_generated_image_urls_from_response_parts } from "./client.ts";

test("response part fallback finds generated images when legacy generated markers are absent", () => {
  const generatedUrl = "https://lh3.googleusercontent.com/gg-dl/example-generated-image";
  const initialCandidate = ["rcid-1", ["image generated successfully"]];
  const imageCandidate = [
    "rcid-1",
    ["image generated successfully"],
    { nestedPayload: [{ media: generatedUrl }] },
  ];
  const responseJson = [
    ["wrb.fr", null, JSON.stringify([null, [], null, null, [initialCandidate]])],
    ["wrb.fr", null, JSON.stringify([null, [], null, null, [imageCandidate]])],
  ];

  assert.equal(initialCandidate[12], undefined);
  assert.equal(
    /http:\/\/googleusercontent\.com\/image_generation_content\/\d+/.test(String(initialCandidate[1]?.[0])),
    false,
  );
  assert.deepEqual(collect_generated_image_urls_from_response_parts(responseJson, 0, 0), [generatedUrl]);
});

import { Endpoint, ErrorCode, Headers, Model } from './constants.js';
import { GemMixin } from './components/gem-mixin.js';
import {
  APIError,
  AuthError,
  GeminiError,
  ImageGenerationError,
  ModelInvalid,
  TemporarilyBlocked,
  TimeoutError,
  UsageLimitExceeded,
} from './exceptions.js';
import { Candidate, Gem, GeneratedImage, ModelOutput, RPCData, WebImage } from './types/index.js';
import {
  extract_json_from_response,
  get_access_token,
  get_nested_value,
  logger,
  parse_file_name,
  rotate_1psidts,
  rotate_tasks,
  fetch_with_timeout,
  sleep,
  upload_file,
  write_cookie_file,
  resolveGeminiWebCookiePath,
} from './utils/index.js';

type InitOptions = {
  timeout?: number;
  auto_close?: boolean;
  close_delay?: number;
  auto_refresh?: boolean;
  refresh_interval?: number;
  verbose?: boolean;
};

type RequestKwargs = RequestInit & { timeout_ms?: number };

const GENERATED_IMAGE_URL_PREFIX = 'https://lh3.googleusercontent.com/gg-dl/';

function normalize_headers(h?: HeadersInit): Record<string, string> {
  if (!h) return {};
  if (Array.isArray(h)) return Object.fromEntries(h.map(([k, v]) => [k, v]));
  if (h instanceof Headers) {
    const out: Record<string, string> = {};
    h.forEach((v, k) => {
      out[k] = v;
    });
    return out;
  }
  return { ...(h as Record<string, string>) };
}

function collect_strings(root: unknown, accept: (s: string) => boolean, limit: number = 20): string[] {
  const out: string[] = [];
  const seen = new Set<string>();
  const stack: unknown[] = [root];

  while (stack.length > 0 && out.length < limit) {
    const v = stack.pop();
    if (typeof v === 'string') {
      if (accept(v) && !seen.has(v)) {
        seen.add(v);
        out.push(v);
      }
      continue;
    }

    if (Array.isArray(v)) {
      for (let i = 0; i < v.length; i++) stack.push(v[i]);
      continue;
    }

    if (v && typeof v === 'object') {
      for (const val of Object.values(v as Record<string, unknown>)) stack.push(val);
    }
  }

  return out;
}

function collect_generated_image_urls(root: unknown, limit: number = 4): string[] {
  return collect_strings(root, (s) => s.startsWith(GENERATED_IMAGE_URL_PREFIX), limit);
}

function parse_response_part_body(part: unknown): unknown[] | null {
  if (!Array.isArray(part)) return null;
  const part_body = get_nested_value<string | null>(part, [2], null);
  if (!part_body) return null;
  try {
    const part_json = JSON.parse(part_body) as unknown;
    return Array.isArray(part_json) ? part_json : null;
  } catch {
    return null;
  }
}

function find_generated_image_part(
  response_json: unknown[],
  body_index: number,
  candidate_index: number,
  limit: number = 4,
): { body: unknown[]; urls: string[] } | null {
  for (let part_index = body_index; part_index < response_json.length; part_index++) {
    const part_json = parse_response_part_body(response_json[part_index]);
    if (!part_json) continue;
    const cand = get_nested_value<unknown>(part_json, [4, candidate_index], null);
    if (!cand) continue;
    const urls = collect_generated_image_urls(cand, limit);
    if (urls.length > 0) return { body: part_json, urls };
  }
  return null;
}

export function collect_generated_image_urls_from_response_parts(
  response_json: unknown[],
  body_index: number,
  candidate_index: number,
  limit: number = 4,
): string[] {
  return find_generated_image_part(response_json, body_index, candidate_index, limit)?.urls ?? [];
}

function push_generated_images(
  generated_images: GeneratedImage[],
  urls: string[],
  proxy: string | null,
  cookies: Record<string, string>,
): void {
  for (const url of urls) {
    generated_images.push(new GeneratedImage(url, '[Generated Image]', '', proxy, cookies));
  }
}

export class GeminiClient extends GemMixin {
  public cookies: Record<string, string> = {};
  public proxy: string | null = null;
  public _running: boolean = false;
  public access_token: string | null = null;
  public timeout: number = 300;
  public auto_close: boolean = false;
  public close_delay: number = 300;
  public auto_refresh: boolean = true;
  public refresh_interval: number = 540;
  public kwargs: RequestInit;

  private close_timer: ReturnType<typeof setTimeout> | null = null;
  private refresh_abort: AbortController | null = null;

  constructor(
    secure_1psid: string | null = null,
    secure_1psidts: string | null = null,
    proxy: string | null = null,
    kwargs: RequestInit = {},
  ) {
    super();
    this.proxy = proxy;
    this.kwargs = kwargs;

    if (secure_1psid) {
      this.cookies['__Secure-1PSID'] = secure_1psid;
      if (secure_1psidts) this.cookies['__Secure-1PSIDTS'] = secure_1psidts;
    }
  }

  async init(
    timeoutOrOpts: number | InitOptions = 300,
    auto_close: boolean = false,
    close_delay: number = 300,
    auto_refresh: boolean = true,
    refresh_interval: number = 540,
    verbose: boolean = true,
  ): Promise<void> {
    const opts: InitOptions =
      typeof timeoutOrOpts === 'object'
        ? timeoutOrOpts
        : { timeout: timeoutOrOpts, auto_close, close_delay, auto_refresh, refresh_interval, verbose };

    const timeout = opts.timeout ?? 300;
    const ac = opts.auto_close ?? false;
    const cd = opts.close_delay ?? 300;
    const ar = opts.auto_refresh ?? true;
    const ri = opts.refresh_interval ?? 540;
    const vb = opts.verbose ?? true;

    try {
      const [token, valid] = await get_access_token(this.cookies, this.proxy, vb);
      this.access_token = token;
      this.cookies = valid;
      this._running = true;

      this.timeout = timeout;
      this.auto_close = ac;
      this.close_delay = cd;
      if (this.auto_close) await this.reset_close_task();

      this.auto_refresh = ar;
      this.refresh_interval = ri;

      const sid = this.cookies['__Secure-1PSID'];
      if (sid) {
        const existing = rotate_tasks.get(sid);
        if (existing && existing instanceof AbortController) existing.abort();
        rotate_tasks.delete(sid);
      }

      if (this.auto_refresh && sid) {
        const ctl = new AbortController();
        this.refresh_abort?.abort();
        this.refresh_abort = ctl;
        rotate_tasks.set(sid, ctl);
        void this.start_auto_refresh(ctl.signal);
      }

      await write_cookie_file(this.cookies, resolveGeminiWebCookiePath(), 'client').catch(() => {});

      if (vb) logger.success('Gemini client initialized successfully.');
    } catch (e) {
      await this.close();
      throw e;
    }
  }

  async close(delay: number = 0): Promise<void> {
    if (delay > 0) await sleep(delay * 1000);
    this._running = false;

    if (this.close_timer) {
      clearTimeout(this.close_timer);
      this.close_timer = null;
    }

    this.refresh_abort?.abort();
    this.refresh_abort = null;

    const sid = this.cookies['__Secure-1PSID'];
    const t = sid ? rotate_tasks.get(sid) : null;
    if (t && t instanceof AbortController) t.abort();
    if (sid) rotate_tasks.delete(sid);
  }

  async reset_close_task(): Promise<void> {
    if (this.close_timer) {
      clearTimeout(this.close_timer);
      this.close_timer = null;
    }

    this.close_timer = setTimeout(() => {
      void this.close(0);
    }, this.close_delay * 1000);
    this.close_timer.unref?.();
  }

  async start_auto_refresh(signal: AbortSignal): Promise<void> {
    while (!signal.aborted) {
      let newTs: string | null = null;
      try {
        newTs = await rotate_1psidts(this.cookies, this.proxy);
      } catch (e) {
        if (e instanceof AuthError) {
          logger.warning('AuthError: Failed to refresh cookies. Auto refresh task canceled.');
          return;
        }
        logger.warning(`Unexpected error while refreshing cookies: ${e instanceof Error ? e.message : String(e)}`);
      }

      if (newTs) {
        this.cookies['__Secure-1PSIDTS'] = newTs;
        await write_cookie_file(this.cookies, resolveGeminiWebCookiePath(), 'refresh').catch(() => {});
        logger.debug('Cookies refreshed. New __Secure-1PSIDTS applied.');
      }

      await sleep(this.refresh_interval * 1000, signal);
    }
  }

  protected async _run<T>(fn: () => Promise<T>, retry: number): Promise<T> {
    try {
      if (!this._running) {
        await this.init({
          timeout: this.timeout,
          auto_close: this.auto_close,
          close_delay: this.close_delay,
          auto_refresh: this.auto_refresh,
          refresh_interval: this.refresh_interval,
          verbose: false,
        });

        if (!this._running) {
          throw new APIError('Client initialization failed.');
        }
      }

      return await fn();
    } catch (e) {
      let r = retry;
      if (e instanceof ImageGenerationError) r = Math.min(1, r);
      if (e instanceof APIError && r > 0) {
        await sleep(1000);
        return await this._run(fn, r - 1);
      }
      throw e;
    }
  }

  async generate_content(
    prompt: string,
    files: string[] | null = null,
    model: Model | string | Record<string, unknown> = Model.UNSPECIFIED,
    gem: Gem | string | null = null,
    chat: ChatSession | null = null,
    kwargs: RequestKwargs = {},
  ): Promise<ModelOutput> {
    return await this._run(async () => {
      if (!prompt) throw new Error('Prompt cannot be empty.');

      let mdl: Model;
      if (typeof model === 'string') mdl = Model.from_name(model);
      else if (model instanceof Model) mdl = model;
      else if (model && typeof model === 'object') mdl = Model.from_dict(model);
      else throw new TypeError(`'model' must be a Model instance, string, or dictionary; got ${typeof model}`);

      const gem_id = gem instanceof Gem ? gem.id : gem;

      if (this.auto_close) await this.reset_close_task();

      if (!this.access_token) throw new APIError('Missing access token.');

      const f = files?.length ? files : null;
      const uploaded =
        f &&
        (await Promise.all(
          f.map(async (p) => [[await upload_file(p, this.proxy)], parse_file_name(p)] as [string[], string]),
        ));

      const first = uploaded ? [prompt, 0, null, uploaded] : [prompt];
      const inner: unknown[] = [first, null, chat ? chat.metadata : null];

      if (gem_id) {
        for (let i = 0; i < 16; i++) inner.push(null);
        inner.push(gem_id);
      }

      const f_req = JSON.stringify([null, JSON.stringify(inner)]);
      const body = new URLSearchParams({ at: this.access_token, 'f.req': f_req }).toString();

      const h0 = { ...Headers.GEMINI, ...mdl.model_header, Cookie: Object.entries(this.cookies).map(([k, v]) => `${k}=${v}`).join('; ') };
      const h1 = { ...h0, ...normalize_headers(kwargs.headers) };

      let res: Response;
      try {
        const timeout_ms = typeof kwargs.timeout_ms === 'number' ? kwargs.timeout_ms : this.timeout * 1000;
        const { timeout_ms: _t, ...rest } = kwargs;
        res = await fetch_with_timeout(Endpoint.GENERATE, {
          method: 'POST',
          headers: h1,
          body,
          redirect: 'follow',
          ...this.kwargs,
          ...rest,
          timeout_ms,
        });
      } catch (e) {
        throw new TimeoutError(
          `Generate content request timed out, please try again. If the problem persists, consider setting a higher 'timeout' value when initializing GeminiClient. (${e instanceof Error ? e.message : String(e)})`,
        );
      }

      if (res.status !== 200) {
        await this.close();
        throw new APIError(`Failed to generate contents. Request failed with status code ${res.status}`);
      }

      const txt = await res.text();
      const response_json = extract_json_from_response(txt);

      let body_json: unknown[] | null = null;
      let body_index = 0;

      try {
        if (!Array.isArray(response_json)) throw new Error('Invalid JSON');
        for (let part_index = 0; part_index < response_json.length; part_index++) {
          const part = response_json[part_index];
          if (!Array.isArray(part)) continue;
          const part_body = get_nested_value<string | null>(part, [2], null);
          if (!part_body) continue;
          try {
            const part_json = JSON.parse(part_body) as unknown[];
            if (get_nested_value(part_json, [4], null)) {
              body_index = part_index;
              body_json = part_json;
              break;
            }
          } catch {}
        }
        if (!body_json) throw new Error('No body');
      } catch {
        await this.close();
        try {
          const code = get_nested_value<number>(response_json, [0, 5, 2, 0, 1, 0], -1);
          if (code === ErrorCode.USAGE_LIMIT_EXCEEDED) {
            throw new UsageLimitExceeded(
              `Failed to generate contents. Usage limit of ${mdl.model_name} model has exceeded. Please try switching to another model.`,
            );
          }
          if (code === ErrorCode.MODEL_INCONSISTENT) {
            throw new ModelInvalid(
              'Failed to generate contents. The specified model is inconsistent with the chat history. Please make sure to pass the same `model` parameter when starting a chat session with previous metadata.',
            );
          }
          if (code === ErrorCode.MODEL_HEADER_INVALID) {
            throw new ModelInvalid(
              'Failed to generate contents. The specified model is not available. Please update gemini_webapi to the latest version. If the error persists and is caused by the package, please report it on GitHub.',
            );
          }
          if (code === ErrorCode.IP_TEMPORARILY_BLOCKED) {
            throw new TemporarilyBlocked(
              'Failed to generate contents. Your IP address is temporarily blocked by Google. Please try using a proxy or waiting for a while.',
            );
          }
        } catch (e) {
          if (e instanceof GeminiError) throw e;
        }

        logger.debug(`Invalid response: ${txt.slice(0, 500)}`);
        throw new APIError('Failed to generate contents. Invalid response data received. Client will try to re-initialize on next request.');
      }

      try {
        const candidate_list = get_nested_value<unknown[]>(body_json, [4], []);
        const out: Candidate[] = [];

        for (let candidate_index = 0; candidate_index < candidate_list.length; candidate_index++) {
          const candidate = candidate_list[candidate_index];
          if (!Array.isArray(candidate)) continue;

          const rcid = get_nested_value<string | null>(candidate, [0], null);
          if (!rcid) continue;

          let text = String(get_nested_value(candidate, [1, 0], ''));
          if (/^http:\/\/googleusercontent\.com\/card_content\/\d+/.test(text)) {
            text = String(get_nested_value(candidate, [22, 0], text));
          }

          const thoughts = get_nested_value<string | null>(candidate, [37, 0, 0], null);

          const web_images: WebImage[] = [];
          for (const w of get_nested_value<unknown[]>(candidate, [12, 1], [])) {
            if (!Array.isArray(w)) continue;
            const url = get_nested_value<string | null>(w, [0, 0, 0], null);
            if (!url) continue;
            web_images.push(new WebImage(url, String(get_nested_value(w, [7, 0], '')), String(get_nested_value(w, [0, 4], '')), this.proxy));
          }

          const generated_images: GeneratedImage[] = [];
          const wants_generated =
            get_nested_value(candidate, [12, 7, 0], null) != null ||
            /http:\/\/googleusercontent\.com\/image_generation_content\/\d+/.test(text) ||
            // Gemini 3 no longer emits the legacy marker in the candidate text; the generated
            // image arrives in a later response part. Detect it from the raw response instead.
            /\/gg-dl\//.test(txt) ||
            /image_generation_content\/\d+/.test(txt);

          if (wants_generated) {
            const image_part = find_generated_image_part(response_json as unknown[], body_index, candidate_index, 1);
            const img_body = image_part?.body ?? null;

            // Not every candidate carries a generated image (e.g. multi-candidate responses).
            // If this candidate has no image part, skip extraction instead of failing the whole call.
            const img_candidate = img_body ? get_nested_value<unknown[]>(img_body, [4, candidate_index], []) : [];
            const finished = get_nested_value<string | null>(img_candidate, [1, 0], null);
            if (finished) {
              text = finished.replace(/http:\/\/googleusercontent\.com\/image_generation_content\/\d+/g, '').trimEnd();
            }

            const gen = get_nested_value<unknown[]>(img_candidate, [12, 7, 0], []);
            for (let img_index = 0; img_index < gen.length; img_index++) {
              const g = gen[img_index];
              if (!Array.isArray(g)) continue;
              const url = get_nested_value<string | null>(g, [0, 3, 3], null);
              if (!url) continue;
              const img_num = get_nested_value<number | null>(g, [3, 6], null);
              const title = img_num ? `[Generated Image ${img_num}]` : '[Generated Image]';
              const alt_list = get_nested_value<unknown[]>(g, [3, 5], []);
              const alt =
                (typeof alt_list[img_index] === 'string' ? (alt_list[img_index] as string) : null) ??
                (typeof alt_list[0] === 'string' ? (alt_list[0] as string) : '') ??
                '';
              generated_images.push(new GeneratedImage(url, title, alt, this.proxy, this.cookies));
            }

            if (generated_images.length === 0) {
              push_generated_images(generated_images, collect_generated_image_urls(img_candidate), this.proxy, this.cookies);
            }
          }

          // Fallback: unconditionally scan all response parts for generated image URLs.
          // The `wants_generated` detection above relies on `candidate[12][7][0]` and an old
          // `googleusercontent.com/image_generation_content/` URL pattern, both of which no
          // longer appear in the current Gemini Web API response format.
          // When Gemini does generate images, their URLs now start with
          // `https://lh3.googleusercontent.com/gg-dl/` and are present somewhere in the
          // response parts — this fallback finds them so images are not silently dropped.
          if (generated_images.length === 0) {
            const urls = collect_generated_image_urls_from_response_parts(response_json as unknown[], body_index, candidate_index);
            push_generated_images(generated_images, urls, this.proxy, this.cookies);
          }

          out.push(new Candidate({ rcid, text, thoughts, web_images, generated_images }));
        }

        if (out.length === 0) {
          throw new GeminiError('Failed to generate contents. No output data found in response.');
        }

        const metadata = get_nested_value<string[]>(body_json, [1], []);
        const output = new ModelOutput({ metadata, candidates: out });

        if (chat instanceof ChatSession) chat.last_output = output;
        return output;
      } catch (e) {
        if (e instanceof GeminiError || e instanceof APIError) throw e;
        throw new APIError('Failed to parse response body. Data structure is invalid.');
      }
    }, 2);
  }

  async generateContent(
    prompt: string,
    files?: string[] | null,
    model?: Model | string | Record<string, unknown>,
    gem?: Gem | string | null,
    chat?: ChatSession | null,
    kwargs?: RequestKwargs,
  ): Promise<ModelOutput> {
    return await this.generate_content(prompt, files ?? null, model ?? Model.UNSPECIFIED, gem ?? null, chat ?? null, kwargs ?? {});
  }

  start_chat(opts?: ConstructorParameters<typeof ChatSession>[1]): ChatSession {
    return new ChatSession(this, opts);
  }

  startChat(opts?: ConstructorParameters<typeof ChatSession>[1]): ChatSession {
    return this.start_chat(opts);
  }

  protected async _batch_execute(payloads: RPCData[], opts: RequestInit = {}): Promise<Response> {
    if (!this.access_token) throw new APIError('Missing access token.');

    const f_req = JSON.stringify([payloads.map((p) => p.serialize())]);
    const body = new URLSearchParams({ at: this.access_token, 'f.req': f_req }).toString();

    const h0 = { ...Headers.GEMINI, Cookie: Object.entries(this.cookies).map(([k, v]) => `${k}=${v}`).join('; ') };
    const h1 = { ...h0, ...normalize_headers(opts.headers) };

    const res = await fetch_with_timeout(Endpoint.BATCH_EXEC, {
      method: 'POST',
      headers: h1,
      body,
      redirect: 'follow',
      ...this.kwargs,
      ...opts,
      timeout_ms: this.timeout * 1000,
    });

    if (res.status !== 200) {
      await this.close();
      throw new APIError(`Batch execution failed with status code ${res.status}`);
    }

    return res;
  }
}

export class ChatSession {
  private __metadata: Array<string | null> = [null, null, null];
  public geminiclient: GeminiClient;
  private _last_output: ModelOutput | null = null;
  public model: Model | string | Record<string, unknown>;
  public gem: Gem | string | null;

  constructor(
    geminiclient: GeminiClient,
    opts: {
      metadata?: Array<string | null>;
      cid?: string | null;
      rid?: string | null;
      rcid?: string | null;
      model?: Model | string | Record<string, unknown>;
      gem?: Gem | string | null;
    } = {},
  ) {
    this.geminiclient = geminiclient;
    this.model = opts.model ?? Model.UNSPECIFIED;
    this.gem = opts.gem ?? null;

    if (opts.metadata) this.metadata = opts.metadata;
    if (opts.cid) this.cid = opts.cid;
    if (opts.rid) this.rid = opts.rid;
    if (opts.rcid) this.rcid = opts.rcid;
  }

  toString(): string {
    return `ChatSession(cid='${this.cid}', rid='${this.rid}', rcid='${this.rcid}')`;
  }

  get last_output(): ModelOutput | null {
    return this._last_output;
  }

  set last_output(v: ModelOutput | null) {
    this._last_output = v;
    if (v) {
      this.metadata = (v.metadata ?? []) as Array<string | null>;
      this.rcid = v.rcid;
    }
  }

  async send_message(prompt: string, files: string[] | null = null, kwargs: RequestKwargs = {}): Promise<ModelOutput> {
    return await this.geminiclient.generate_content(prompt, files, this.model, this.gem, this, kwargs);
  }

  async sendMessage(prompt: string, files?: string[] | null, kwargs?: RequestKwargs): Promise<ModelOutput> {
    return await this.send_message(prompt, files ?? null, kwargs ?? {});
  }

  choose_candidate(index: number): ModelOutput {
    if (!this.last_output) throw new Error('No previous output data found in this chat session.');
    if (index >= this.last_output.candidates.length) {
      throw new Error(`Index ${index} exceeds the number of candidates in last model output.`);
    }
    this.last_output.chosen = index;
    this.rcid = this.last_output.rcid;
    return this.last_output;
  }

  chooseCandidate(index: number): ModelOutput {
    return this.choose_candidate(index);
  }

  get metadata(): Array<string | null> {
    return this.__metadata;
  }

  set metadata(v: Array<string | null>) {
    if (v.length > 3) throw new Error('metadata cannot exceed 3 elements');
    this.__metadata = [null, null, null];
    for (let i = 0; i < v.length; i++) this.__metadata[i] = v[i] ?? null;
  }

  get cid(): string | null {
    return this.__metadata[0];
  }

  set cid(v: string | null) {
    this.__metadata[0] = v;
  }

  get rid(): string | null {
    return this.__metadata[1];
  }

  set rid(v: string | null) {
    this.__metadata[1] = v;
  }

  get rcid(): string | null {
    return this.__metadata[2];
  }

  set rcid(v: string | null) {
    this.__metadata[2] = v;
  }
}

import { GRPC } from '../constants.js';
import { APIError } from '../exceptions.js';
import { Gem, GemJar, RPCData } from '../types/index.js';
import { logger } from '../utils/logger.js';
import { extract_json_from_response, get_nested_value } from '../utils/parsing.js';

export abstract class GemMixin {
  protected _gems: GemJar | null = null;

  protected abstract _run<T>(fn: () => Promise<T>, retry: number): Promise<T>;
  protected abstract _batch_execute(payloads: RPCData[], opts?: RequestInit): Promise<Response>;
  protected abstract close(delay?: number): Promise<void>;

  get gems(): GemJar {
    if (this._gems == null) {
      throw new Error(
        'Gems not fetched yet. Call `GeminiClient.fetch_gems()` method to fetch gems from gemini.google.com.',
      );
    }
    return this._gems;
  }

  async fetch_gems(include_hidden: boolean = false, opts?: RequestInit): Promise<GemJar> {
    return await this._run(async () => {
      const res = await this._batch_execute(
        [
          new RPCData(GRPC.LIST_GEMS, include_hidden ? '[4]' : '[3]', 'system'),
          new RPCData(GRPC.LIST_GEMS, '[2]', 'custom'),
        ],
        opts,
      );

      let response_json: unknown;
      try {
        response_json = extract_json_from_response(await res.text());
        if (!Array.isArray(response_json)) throw new Error('Invalid response');
      } catch {
        await this.close();
        throw new APIError('Failed to fetch gems. Invalid response data received. Client will try to re-initialize on next request.');
      }

      let predefined: unknown[] = [];
      let custom: unknown[] = [];

      try {
        for (const part of response_json as unknown[]) {
          if (!Array.isArray(part)) continue;
          const ident = part[part.length - 1];
          const body = get_nested_value<string | null>(part, [2], null);
          if (!body) continue;

          if (ident === 'system') {
            const parsed = JSON.parse(body) as unknown[];
            predefined = (Array.isArray(parsed) ? (parsed[2] as unknown[]) : []) ?? [];
          } else if (ident === 'custom') {
            const parsed = JSON.parse(body) as unknown[] | null;
            if (parsed) custom = (parsed[2] as unknown[]) ?? [];
          }
        }

        if (predefined.length === 0 && custom.length === 0) throw new Error('No gems');
      } catch {
        await this.close();
        logger.debug('Invalid response while parsing gems');
        throw new APIError('Failed to fetch gems. Invalid response data received. Client will try to re-initialize on next request.');
      }

      const entries: [string, Gem][] = [];

      for (const gem of predefined) {
        if (!Array.isArray(gem)) continue;
        const id = String(get_nested_value(gem, [0], ''));
        if (!id) continue;
        entries.push([
          id,
          new Gem(
            id,
            String(get_nested_value(gem, [1, 0], '')),
            get_nested_value<string | null>(gem, [1, 1], null),
            get_nested_value<string | null>(gem, [2, 0], null),
            true,
          ),
        ]);
      }

      for (const gem of custom) {
        if (!Array.isArray(gem)) continue;
        const id = String(get_nested_value(gem, [0], ''));
        if (!id) continue;
        entries.push([
          id,
          new Gem(
            id,
            String(get_nested_value(gem, [1, 0], '')),
            get_nested_value<string | null>(gem, [1, 1], null),
            get_nested_value<string | null>(gem, [2, 0], null),
            false,
          ),
        ]);
      }

      this._gems = new GemJar(entries);
      return this._gems;
    }, 2);
  }

  async create_gem(name: string, prompt: string, description: string = ''): Promise<Gem> {
    return await this._run(async () => {
      const payload = JSON.stringify([
        [
          name,
          description,
          prompt,
          null,
          null,
          null,
          null,
          null,
          0,
          null,
          1,
          null,
          null,
          null,
          [],
        ],
      ]);

      const res = await this._batch_execute([new RPCData(GRPC.CREATE_GEM, payload)]);
      try {
        const response_json = extract_json_from_response(await res.text()) as unknown[];
        const gem_id = JSON.parse(String((response_json[0] as unknown[])[2]))[0] as string;
        return new Gem(gem_id, name, description, prompt, false);
      } catch {
        await this.close();
        throw new APIError('Failed to create gem. Invalid response data received. Client will try to re-initialize on next request.');
      }
    }, 2);
  }

  async update_gem(gem: Gem | string, name: string, prompt: string, description: string = ''): Promise<Gem> {
    return await this._run(async () => {
      const gem_id = typeof gem === 'string' ? gem : gem.id;
      const payload = JSON.stringify([
        gem_id,
        [
          name,
          description,
          prompt,
          null,
          null,
          null,
          null,
          null,
          0,
          null,
          1,
          null,
          null,
          null,
          [],
          0,
        ],
      ]);

      await this._batch_execute([new RPCData(GRPC.UPDATE_GEM, payload)]);
      return new Gem(gem_id, name, description, prompt, false);
    }, 2);
  }

  async delete_gem(gem: Gem | string, opts?: RequestInit): Promise<void> {
    return await this._run(async () => {
      const gem_id = typeof gem === 'string' ? gem : gem.id;
      const payload = JSON.stringify([gem_id]);
      await this._batch_execute([new RPCData(GRPC.DELETE_GEM, payload)], opts);
    }, 2);
  }
}

export { GemMixin } from './gem-mixin.js';

export const Endpoint = {
  GOOGLE: 'https://www.google.com',
  INIT: 'https://gemini.google.com/app',
  GENERATE:
    'https://gemini.google.com/_/BardChatUi/data/assistant.lamda.BardFrontendService/StreamGenerate',
  ROTATE_COOKIES: 'https://accounts.google.com/RotateCookies',
  UPLOAD: 'https://content-push.googleapis.com/upload',
  BATCH_EXEC: 'https://gemini.google.com/_/BardChatUi/data/batchexecute',
} as const;

export const GRPC = {
  LIST_CHATS: 'MaZiqc',
  READ_CHAT: 'hNvQHb',
  LIST_GEMS: 'CNgdBe',
  CREATE_GEM: 'oMH3Zd',
  UPDATE_GEM: 'kHv0Vd',
  DELETE_GEM: 'UXcSJb',
} as const;

export const Headers = {
  GEMINI: {
    'Content-Type': 'application/x-www-form-urlencoded;charset=utf-8',
    Host: 'gemini.google.com',
    Origin: 'https://gemini.google.com',
    Referer: 'https://gemini.google.com/',
    'User-Agent':
      'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36',
    'X-Same-Domain': '1',
  },
  ROTATE_COOKIES: {
    'Content-Type': 'application/json',
  },
  UPLOAD: {
    'Push-ID': 'feeds/mcudyrk2a4khkz',
  },
} as const;

export const ErrorCode = {
  TEMPORARY_ERROR_1013: 1013,
  USAGE_LIMIT_EXCEEDED: 1037,
  MODEL_INCONSISTENT: 1050,
  MODEL_HEADER_INVALID: 1052,
  IP_TEMPORARILY_BLOCKED: 1060,
} as const;

export class Model {
  static readonly UNSPECIFIED = new Model('unspecified', {}, false);
  static readonly G_3_0_PRO = new Model(
    'gemini-3.0-pro',
    { 'x-goog-ext-525001261-jspb': '[1,null,null,null,"9d8ca3786ebdfbea",null,null,0,[4],null,null,1]' },
    false,
  );
  static readonly G_3_0_FLASH = new Model(
    'gemini-3.0-flash',
    { 'x-goog-ext-525001261-jspb': '[1,null,null,null,"fbb127bbb056c959",null,null,0,[4],null,null,1]' },
    false,
  );
  static readonly G_3_0_FLASH_THINKING = new Model(
    'gemini-3.0-flash-thinking',
    { 'x-goog-ext-525001261-jspb': '[1,null,null,null,"5bf011840784117a",null,null,0,[4],null,null,1]' },
    false,
  );
  static readonly G_3_1_PRO_PREVIEW = new Model(
    'gemini-3.1-pro-preview',
    {},
    false,
  );

  constructor(
    public readonly model_name: string,
    public readonly model_header: Record<string, string>,
    public readonly advanced_only: boolean,
  ) {}

  static from_name(name: string): Model {
    for (const model of [Model.UNSPECIFIED, Model.G_3_0_PRO, Model.G_3_0_FLASH, Model.G_3_0_FLASH_THINKING, Model.G_3_1_PRO_PREVIEW]) {
      if (model.model_name === name) return model;
    }

    throw new Error(
      `Unknown model name: ${name}. Available models: ${[Model.UNSPECIFIED, Model.G_3_0_PRO, Model.G_3_0_FLASH, Model.G_3_0_FLASH_THINKING, Model.G_3_1_PRO_PREVIEW]
        .map((m) => m.model_name)
        .join(', ')}`,
    );
  }

  static from_dict(model_dict: { model_name?: unknown; model_header?: unknown }): Model {
    if (!model_dict || typeof model_dict !== 'object') {
      throw new Error("When passing a custom model as a dictionary, 'model_name' and 'model_header' keys must be provided.");
    }

    if (!('model_name' in model_dict) || !('model_header' in model_dict)) {
      throw new Error("When passing a custom model as a dictionary, 'model_name' and 'model_header' keys must be provided.");
    }

    if (typeof model_dict.model_name !== 'string' || !model_dict.model_name.trim()) {
      throw new Error("When passing a custom model as a dictionary, 'model_name' must be a non-empty string.");
    }

    if (!model_dict.model_header || typeof model_dict.model_header !== 'object') {
      throw new Error("When passing a custom model as a dictionary, 'model_header' must be a dictionary containing valid header strings.");
    }

    const header: Record<string, string> = {};
    for (const [k, v] of Object.entries(model_dict.model_header as Record<string, unknown>)) {
      if (typeof v === 'string') header[k] = v;
    }

    return new Model(model_dict.model_name, header, false);
  }
}

export class AuthError extends Error {
  constructor(message = 'AuthError') {
    super(message);
    this.name = 'AuthError';
  }
}

export class APIError extends Error {
  constructor(message = 'APIError') {
    super(message);
    this.name = 'APIError';
  }
}

export class ImageGenerationError extends APIError {
  constructor(message = 'ImageGenerationError') {
    super(message);
    this.name = 'ImageGenerationError';
  }
}

export class GeminiError extends Error {
  constructor(message = 'GeminiError') {
    super(message);
    this.name = 'GeminiError';
  }
}

export class TimeoutError extends GeminiError {
  constructor(message = 'TimeoutError') {
    super(message);
    this.name = 'TimeoutError';
  }
}

export class UsageLimitExceeded extends GeminiError {
  constructor(message = 'UsageLimitExceeded') {
    super(message);
    this.name = 'UsageLimitExceeded';
  }
}

export class ModelInvalid extends GeminiError {
  constructor(message = 'ModelInvalid') {
    super(message);
    this.name = 'ModelInvalid';
  }
}

export class TemporarilyBlocked extends GeminiError {
  constructor(message = 'TemporarilyBlocked') {
    super(message);
    this.name = 'TemporarilyBlocked';
  }
}

export { GeminiClient, ChatSession } from './client.js';

export * from './exceptions.js';
export * from './types/index.js';
export * from './constants.js';
export { logger, set_log_level, setLogLevel } from './utils/logger.js';
export * as utils from './utils/index.js';

import { GeneratedImage, type Image, WebImage } from './image.js';

function decode_html(s: string | null | undefined): string | null | undefined {
  if (s == null) return s;
  return s
    .replace(/&/g, '&')
    .replace(/&lt;/g, '<')
    .replace(/&gt;/g, '>')
    .replace(/&quot;/g, '"')
    .replace(/&apos;/g, "'")
    .replace(/&#39;/g, "'")
    .replace(/&nbsp;/g, ' ')
    .replace(/&#x([0-9a-fA-F]+);/g, (_, hex) => String.fromCodePoint(parseInt(hex, 16)))
    .replace(/&#(\d+);/g, (_, dec) => String.fromCodePoint(parseInt(dec, 10)));
}

export class Candidate {
  public rcid: string;
  public text: string;
  public thoughts: string | null;
  public web_images: WebImage[];
  public generated_images: GeneratedImage[];

  constructor(params: {
    rcid: string;
    text: string;
    thoughts?: string | null;
    web_images?: WebImage[];
    generated_images?: GeneratedImage[];
  }) {
    this.rcid = params.rcid;
    this.text = decode_html(params.text) ?? '';
    this.thoughts = decode_html(params.thoughts) ?? null;
    this.web_images = params.web_images ?? [];
    this.generated_images = params.generated_images ?? [];
  }

  toString(): string {
    return this.text;
  }

  get images(): Image[] {
    return [...this.web_images, ...this.generated_images];
  }
}

export class Gem {
  constructor(
    public id: string,
    public name: string,
    public description: string | null,
    public prompt: string | null,
    public predefined: boolean,
  ) {}

  toString(): string {
    return `Gem(id='${this.id}', name='${this.name}', description='${this.description}', prompt='${this.prompt}', predefined=${this.predefined})`;
  }
}

export class GemJar implements Iterable<Gem> {
  private m = new Map<string, Gem>();

  constructor(entries?: Iterable<[string, Gem]>) {
    if (entries) for (const [id, gem] of entries) this.m.set(id, gem);
  }

  [Symbol.iterator](): Iterator<Gem> {
    return this.m.values();
  }

  entries(): IterableIterator<[string, Gem]> {
    return this.m.entries();
  }

  values(): IterableIterator<Gem> {
    return this.m.values();
  }

  has(id: string): boolean {
    return this.m.has(id);
  }

  set(id: string, gem: Gem): this {
    this.m.set(id, gem);
    return this;
  }

  get(id?: string | null, name?: string | null, def: Gem | null = null): Gem | null {
    if (id == null && name == null) {
      throw new Error('At least one of gem id or name must be provided.');
    }

    if (id != null) {
      const g = this.m.get(id) ?? null;
      if (!g) return def;
      if (name != null) return g.name === name ? g : def;
      return g;
    }

    if (name != null) {
      for (const g of this.m.values()) {
        if (g.name === name) return g;
      }
      return def;
    }

    return def;
  }

  filter(predefined: boolean | null = null, name: string | null = null): GemJar {
    const out: [string, Gem][] = [];
    for (const [id, gem] of this.m.entries()) {
      if (predefined != null && gem.predefined !== predefined) continue;
      if (name != null && gem.name !== name) continue;
      out.push([id, gem]);
    }
    return new GemJar(out);
  }
}

export class RPCData {
  constructor(
    public rpcid: string,
    public payload: string,
    public identifier: string = 'generic',
  ) {}

  toString(): string {
    return `GRPC(rpcid='${this.rpcid}', payload='${this.payload}', identifier='${this.identifier}')`;
  }

  serialize(): unknown[] {
    return [this.rpcid, this.payload, null, this.identifier];
  }
}

import path from 'node:path';
import { mkdir, writeFile } from 'node:fs/promises';

import { logger } from '../utils/logger.js';
import { cookie_header, fetch_with_timeout } from '../utils/http.js';

export class Image {
  constructor(
    public url: string,
    public title = '[Image]',
    public alt = '',
    public proxy: string | null = null,
  ) {}

  toString(): string {
    const u = this.url.length <= 20 ? this.url : `${this.url.slice(0, 8)}...${this.url.slice(-12)}`;
    return `Image(title='${this.title}', alt='${this.alt}', url='${u}')`;
  }

  async save(
    p: string = 'temp',
    filename: string | null = null,
    cookies: Record<string, string> | null = null,
    verbose: boolean = false,
    skip_invalid_filename: boolean = false,
  ): Promise<string | null> {
    filename = filename ?? this.url.split('/').pop()?.split('?')[0] ?? 'image';
    const m = filename.match(/^(.*\.\w+)/);
    if (m) filename = m[1]!;
    else {
      if (verbose) logger.warning(`Invalid filename: ${filename}`);
      if (skip_invalid_filename) return null;
    }

    const headers: Record<string, string> = {
      'User-Agent': 'Mozilla/5.0',
      Accept: 'image/avif,image/webp,image/apng,image/*,*/*;q=0.8',
      Referer: 'https://gemini.google.com/',
    };
    if (cookies) headers.Cookie = cookie_header(cookies);

    let url = this.url;
    let res: Response | null = null;
    for (let i = 0; i < 10; i++) {
      res = await fetch_with_timeout(url, {
        method: 'GET',
        headers,
        redirect: 'manual',
        timeout_ms: 30_000,
      });

      if (res.status >= 300 && res.status < 400) {
        const loc = res.headers.get('location');
        if (!loc) break;
        url = new URL(loc, url).toString();
        continue;
      }

      break;
    }

    if (!res) throw new Error('Image download failed: no response');

    if (!res.ok) {
      throw new Error(`Error downloading image: ${res.status} ${res.statusText}`);
    }

    const ct = res.headers.get('content-type');
    if (ct && !ct.includes('image')) {
      logger.warning(`Content type of ${filename} is not image, but ${ct}.`);
    }

    const dir = path.resolve(p);
    await mkdir(dir, { recursive: true });

    const dest = path.join(dir, filename);
    const buf = Buffer.from(await res.arrayBuffer());
    await writeFile(dest, buf);

    if (verbose) logger.info(`Image saved as ${dest}`);
    return dest;
  }
}

export class WebImage extends Image {}

export class GeneratedImage extends Image {
  constructor(
    url: string,
    title: string,
    alt: string,
    proxy: string | null,
    public cookies: Record<string, string>,
  ) {
    super(url, title, alt, proxy);
    if (!cookies || Object.keys(cookies).length === 0) {
      throw new Error('GeneratedImage is designed to be initialized with same cookies as GeminiClient.');
    }
  }

  async save(
    p: string = 'temp',
    filename: string | null = null,
    cookies: Record<string, string> | null = null,
    verbose: boolean = false,
    skip_invalid_filename: boolean = false,
    full_size: boolean = true,
  ): Promise<string | null> {
    const u = full_size ? `${this.url}=s2048` : this.url;
    const f = filename ?? `${new Date().toISOString().replace(/[-:.TZ]/g, '').slice(0, 14)}_${u.slice(-10)}.png`;
    const img = new Image(u, this.title, this.alt, this.proxy);
    return await img.save(p, f, cookies ?? this.cookies, verbose, skip_invalid_filename);
  }
}