1import path from 'node:path';
2import { mkdir, writeFile } from 'node:fs/promises';
3 
4import { logger } from '../utils/logger.js';
5import { cookie_header, fetch_with_timeout } from '../utils/http.js';
6 
7export class Image {
8  constructor(
9    public url: string,
10    public title = '[Image]',
11    public alt = '',
12    public proxy: string | null = null,
13  ) {}
14 
15  toString(): string {
16    const u = this.url.length <= 20 ? this.url : `${this.url.slice(0, 8)}...${this.url.slice(-12)}`;
17    return `Image(title='${this.title}', alt='${this.alt}', url='${u}')`;
18  }
19 
20  async save(
21    p: string = 'temp',
22    filename: string | null = null,
23    cookies: Record<string, string> | null = null,
24    verbose: boolean = false,
25    skip_invalid_filename: boolean = false,
26  ): Promise<string | null> {
27    filename = filename ?? this.url.split('/').pop()?.split('?')[0] ?? 'image';
28    const m = filename.match(/^(.*\.\w+)/);
29    if (m) filename = m[1]!;
30    else {
31      if (verbose) logger.warning(`Invalid filename: ${filename}`);
32      if (skip_invalid_filename) return null;
33    }
34 
35    const headers: Record<string, string> = {
36      'User-Agent': 'Mozilla/5.0',
37      Accept: 'image/avif,image/webp,image/apng,image/*,*/*;q=0.8',
38      Referer: 'https://gemini.google.com/',
39    };
40    if (cookies) headers.Cookie = cookie_header(cookies);
41 
42    let url = this.url;
43    let res: Response | null = null;
44    for (let i = 0; i < 10; i++) {
45      res = await fetch_with_timeout(url, {
46        method: 'GET',
47        headers,
48        redirect: 'manual',
49        timeout_ms: 30_000,
50      });
51 
52      if (res.status >= 300 && res.status < 400) {
53        const loc = res.headers.get('location');
54        if (!loc) break;
55        url = new URL(loc, url).toString();
56        continue;
57      }
58 
59      break;
60    }
61 
62    if (!res) throw new Error('Image download failed: no response');
63 
64    if (!res.ok) {
65      throw new Error(`Error downloading image: ${res.status} ${res.statusText}`);
66    }
67 
68    const ct = res.headers.get('content-type');
69    if (ct && !ct.includes('image')) {
70      logger.warning(`Content type of ${filename} is not image, but ${ct}.`);
71    }
72 
73    const dir = path.resolve(p);
74    await mkdir(dir, { recursive: true });
75 
76    const dest = path.join(dir, filename);
77    const buf = Buffer.from(await res.arrayBuffer());
78    await writeFile(dest, buf);
79 
80    if (verbose) logger.info(`Image saved as ${dest}`);
81    return dest;
82  }
83}
84 
85export class WebImage extends Image {}
86 
87export class GeneratedImage extends Image {
88  constructor(
89    url: string,
90    title: string,
91    alt: string,
92    proxy: string | null,
93    public cookies: Record<string, string>,
94  ) {
95    super(url, title, alt, proxy);
96    if (!cookies || Object.keys(cookies).length === 0) {
97      throw new Error('GeneratedImage is designed to be initialized with same cookies as GeminiClient.');
98    }
99  }
100 
101  async save(
102    p: string = 'temp',
103    filename: string | null = null,
104    cookies: Record<string, string> | null = null,
105    verbose: boolean = false,
106    skip_invalid_filename: boolean = false,
107    full_size: boolean = true,
108  ): Promise<string | null> {
109    const u = full_size ? `${this.url}=s2048` : this.url;
110    const f = filename ?? `${new Date().toISOString().replace(/[-:.TZ]/g, '').slice(0, 14)}_${u.slice(-10)}.png`;
111    const img = new Image(u, this.title, this.alt, this.proxy);
112    return await img.save(p, f, cookies ?? this.cookies, verbose, skip_invalid_filename);
113  }
114}
115