web.ts

import type { Device, ScreenInfo, DeviceConfig, SampleConfig } from '.';
import puppeteer, { type Browser, type Page, type Cookie } from 'puppeteer';
import { PuppeteerAgent } from '@midscene/web/puppeteer';
import { overrideAIConfig } from '@midscene/web/puppeteer';
import type { Sample } from '../blackbox';

interface Snapshot {
  url: string;
  html: string;
  cookies: Cookie[];
  localStorage: Record<string, string>;
  sessionStorage: Record<string, string>;
  scroll: { x: number; y: number };
}

function canonicalizeUrl(inputUrl: string): string {
  try {
    const url = new URL(inputUrl);
    // Remove fragment
    url.hash = '';
    // Normalize pathname (e.g., remove duplicate slashes)
    url.pathname = url.pathname.replace(/\/{2,}/g, '/');
    // Remove trailing '?' if no search parameters exist
    if (!url.searchParams.toString()) {
      url.search = '';
    }
    // Optionally, remove trailing slash from pathname (except root)
    if (url.pathname.length > 1 && url.pathname.endsWith('/')) {
      url.pathname = url.pathname.slice(0, -1);
    }

    return url.toString();
  } catch (error) {
    throw new Error(`Invalid URL: ${inputUrl}`);
  }
}

export class WebBrowser implements Device {
  private browser!: Browser;
  private page!: Page;
  private agent!: PuppeteerAgent;
  private readonly snapshots: Map<string, Snapshot> = new Map();

  async start(config: DeviceConfig) {
    const args = [];
    if (config.proxyUrl && config.proxyCertFingerprint) {
      args.push(
        `--proxy-server=${config.proxyUrl}`,
        // '--ignore-certificate-errors',
        '--ignore-certificate-errors-spki-list=' + config.proxyCertFingerprint
      );
    }

    this.browser = await puppeteer.launch({
      headless: false,
      args
    });
  }

  async stop() {
    if (this.browser) {
      await this.browser.close();
    }
  }

  async setup(sample: Sample) {
    this.page = await this.browser.newPage();
    overrideAIConfig({
      OPENAI_API_KEY: process.env.OPENAI_API_KEY,
      OPENAI_BASE_URL: process.env.OPENAI_BASE_URL,
      MIDSCENE_MODEL_NAME: process.env.MIDSCENE_MODEL_NAME,
      MIDSCENE_USE_QWEN_VL: process.env.MIDSCENE_USE_QWEN_VL
    });
    this.agent = new PuppeteerAgent(this.page, {
      forceSameTabNavigation: true,
      // generateReport: false,
      aiActionContext: undefined
    });
    await this.page.setViewport({
      width: 1280,
      height: 800,
      deviceScaleFactor: 1
    });
  }

  async cleanup() {
    if (this.page) {
      await this.page.close();
    }
  }

  async open(sample: Sample) {
    await this.page.goto(sample.target, {
      waitUntil: 'networkidle2'
    });
  }

  async extractUI(): Promise<ScreenInfo> {
    const prompt = `Identify all interactive elements on the page. For each, provide:
- Type (e.g., button, link, input, dropdown, etc.)
- Name (a short, human-readable label)
- Description (a clear explanation of its function or purpose)
If there are many similar elements (e.g., category links, product cards), group them and describe their shared purpose in a single entry.;
The structure of your response should be {type, name, description}[]`;
    const elements = await this.agent.aiQuery(prompt);
    const screenshotBase64 = await this.page.screenshot({
      encoding: 'base64'
    });

    return {
      id: this.page.url(),
      screenshotBase64: `data:image/png;base64,${screenshotBase64}`,
      elements
    };
  }

  async takeSnapshot(id: string): Promise<void> {
    if (this.snapshots.has(id)) return;

    const snapshot: Snapshot = {
      url: this.page.url(),
      html: await this.page.content(),
      cookies: await this.browser.cookies(),
      localStorage: await this.page.evaluate(() =>
        Object.fromEntries(Object.entries(localStorage))
      ),
      sessionStorage: await this.page.evaluate(() =>
        Object.fromEntries(Object.entries(sessionStorage))
      ),
      scroll: await this.page.evaluate(() => ({ x: window.scrollX, y: window.scrollY }))
    };

    this.snapshots.set(id, snapshot);
  }

  async restoreSnapshot(id: string): Promise<void> {
    const snap = this.snapshots.get(id);
    if (!snap) throw new Error('Snapshot not found');

    await this.page.goto(snap.url);

    // Set cookies
    await this.browser.deleteCookie(...snap.cookies);
    await this.browser.setCookie(...snap.cookies);

    // Restore localStorage
    await this.page.evaluate(() => localStorage.clear());
    await this.page.evaluate((data) => {
      for (const [k, v] of Object.entries(data)) localStorage.setItem(k, v);
    }, snap.localStorage);

    // Restore sessionStorage
    await this.page.evaluate(() => sessionStorage.clear());
    await this.page.evaluate((data) => {
      for (const [k, v] of Object.entries(data)) sessionStorage.setItem(k, v);
    }, snap.sessionStorage);

    // Refresh to apply cookies
    await this.page.reload({ waitUntil: 'networkidle2' });
    await this.page.evaluate(({ x, y }) => window.scrollTo(x, y), snap.scroll);
  }

  async removeSnapshot(id: string): Promise<void> {
    this.snapshots.delete(id);
  }

  async takeScreenshot(): Promise<Uint8Array> {
    return await this.page.screenshot();
  }

  async executeAction(action: string): Promise<void> {
    await this.agent.aiAction(action);
  }

  // Getters and setters

  async getUrl(): Promise<string> {
    return this.page.url();
  }
}