import path from "node:path";
import { ensureMediaDir, saveMediaBuffer } from "../../media/store.js";
import { captureScreenshot, snapshotAria } from "../cdp.js";
import {
  evaluateChromeMcpScript,
  navigateChromeMcpPage,
  takeChromeMcpScreenshot,
  takeChromeMcpSnapshot,
} from "../chrome-mcp.js";
import {
  buildAiSnapshotFromChromeMcpSnapshot,
  flattenChromeMcpSnapshotToAriaNodes,
} from "../chrome-mcp.snapshot.js";
import {
  assertBrowserNavigationAllowed,
  assertBrowserNavigationResultAllowed,
} from "../navigation-guard.js";
import { withBrowserNavigationPolicy } from "../navigation-guard.js";
import { getBrowserProfileCapabilities } from "../profile-capabilities.js";
import {
  DEFAULT_BROWSER_SCREENSHOT_MAX_BYTES,
  DEFAULT_BROWSER_SCREENSHOT_MAX_SIDE,
  normalizeBrowserScreenshot,
} from "../screenshot.js";
import type { BrowserRouteContext } from "../server-context.js";
import {
  getPwAiModule,
  handleRouteError,
  readBody,
  requirePwAi,
  resolveProfileContext,
  withPlaywrightRouteContext,
  withRouteTabContext,
} from "./agent.shared.js";
import {
  resolveSnapshotPlan,
  shouldUsePlaywrightForAriaSnapshot,
  shouldUsePlaywrightForScreenshot,
} from "./agent.snapshot.plan.js";
import type { BrowserResponse, BrowserRouteRegistrar } from "./types.js";
import { jsonError, toBoolean, toStringOrEmpty } from "./utils.js";

const CHROME_MCP_OVERLAY_ATTR = "data-openclaw-mcp-overlay";

async function clearChromeMcpOverlay(params: {
  profileName: string;
  userDataDir?: string;
  targetId: string;
}): Promise<void> {
  await evaluateChromeMcpScript({
    profileName: params.profileName,
    userDataDir: params.userDataDir,
    targetId: params.targetId,
    fn: `() => {
      document.querySelectorAll("[${CHROME_MCP_OVERLAY_ATTR}]").forEach((node) => node.remove());
      return true;
    }`,
  }).catch(() => {});
}

async function renderChromeMcpLabels(params: {
  profileName: string;
  userDataDir?: string;
  targetId: string;
  refs: string[];
}): Promise<{ labels: number; skipped: number }> {
  const refList = JSON.stringify(params.refs);
  const result = await evaluateChromeMcpScript({
    profileName: params.profileName,
    userDataDir: params.userDataDir,
    targetId: params.targetId,
    args: params.refs,
    fn: `(...elements) => {
      const refs = ${refList};
      document.querySelectorAll("[${CHROME_MCP_OVERLAY_ATTR}]").forEach((node) => node.remove());
      const root = document.createElement("div");
      root.setAttribute("${CHROME_MCP_OVERLAY_ATTR}", "labels");
      root.style.position = "fixed";
      root.style.inset = "0";
      root.style.pointerEvents = "none";
      root.style.zIndex = "2147483647";
      let labels = 0;
      let skipped = 0;
      elements.forEach((el, index) => {
        if (!(el instanceof Element)) {
          skipped += 1;
          return;
        }
        const rect = el.getBoundingClientRect();
        if (rect.width <= 0 && rect.height <= 0) {
          skipped += 1;
          return;
        }
        labels += 1;
        const badge = document.createElement("div");
        badge.setAttribute("${CHROME_MCP_OVERLAY_ATTR}", "label");
        badge.textContent = refs[index] || String(labels);
        badge.style.position = "fixed";
        badge.style.left = \`\${Math.max(0, rect.left)}px\`;
        badge.style.top = \`\${Math.max(0, rect.top)}px\`;
        badge.style.transform = "translateY(-100%)";
        badge.style.padding = "2px 6px";
        badge.style.borderRadius = "999px";
        badge.style.background = "#FF4500";
        badge.style.color = "#fff";
        badge.style.font = "600 12px ui-monospace, SFMono-Regular, Menlo, monospace";
        badge.style.boxShadow = "0 2px 6px rgba(0,0,0,0.35)";
        badge.style.whiteSpace = "nowrap";
        root.appendChild(badge);
      });
      document.documentElement.appendChild(root);
      return { labels, skipped };
    }`,
  });
  const labels =
    result &&
    typeof result === "object" &&
    typeof (result as { labels?: unknown }).labels === "number"
      ? (result as { labels: number }).labels
      : 0;
  const skipped =
    result &&
    typeof result === "object" &&
    typeof (result as { skipped?: unknown }).skipped === "number"
      ? (result as { skipped: number }).skipped
      : 0;
  return { labels, skipped };
}

async function saveNormalizedScreenshotResponse(params: {
  res: BrowserResponse;
  buffer: Buffer;
  type: "png" | "jpeg";
  targetId: string;
  url: string;
}) {
  const normalized = await normalizeBrowserScreenshot(params.buffer, {
    maxSide: DEFAULT_BROWSER_SCREENSHOT_MAX_SIDE,
    maxBytes: DEFAULT_BROWSER_SCREENSHOT_MAX_BYTES,
  });
  await saveBrowserMediaResponse({
    res: params.res,
    buffer: normalized.buffer,
    contentType: normalized.contentType ?? `image/${params.type}`,
    maxBytes: DEFAULT_BROWSER_SCREENSHOT_MAX_BYTES,
    targetId: params.targetId,
    url: params.url,
  });
}

async function saveBrowserMediaResponse(params: {
  res: BrowserResponse;
  buffer: Buffer;
  contentType: string;
  maxBytes: number;
  targetId: string;
  url: string;
}) {
  await ensureMediaDir();
  const saved = await saveMediaBuffer(
    params.buffer,
    params.contentType,
    "browser",
    params.maxBytes,
  );
  params.res.json({
    ok: true,
    path: path.resolve(saved.path),
    targetId: params.targetId,
    url: params.url,
  });
}

/** Resolve the correct targetId after a navigation that may trigger a renderer swap. */
export async function resolveTargetIdAfterNavigate(opts: {
  oldTargetId: string;
  navigatedUrl: string;
  listTabs: () => Promise<Array<{ targetId: string; url: string }>>;
}): Promise<string> {
  let currentTargetId = opts.oldTargetId;
  try {
    const pickReplacement = (
      tabs: Array<{ targetId: string; url: string }>,
      options?: { allowSingleTabFallback?: boolean },
    ) => {
      if (tabs.some((tab) => tab.targetId === opts.oldTargetId)) {
        return opts.oldTargetId;
      }
      const byUrl = tabs.filter((tab) => tab.url === opts.navigatedUrl);
      if (byUrl.length === 1) {
        return byUrl[0]?.targetId ?? opts.oldTargetId;
      }
      const uniqueReplacement = byUrl.filter((tab) => tab.targetId !== opts.oldTargetId);
      if (uniqueReplacement.length === 1) {
        return uniqueReplacement[0]?.targetId ?? opts.oldTargetId;
      }
      if (options?.allowSingleTabFallback && tabs.length === 1) {
        return tabs[0]?.targetId ?? opts.oldTargetId;
      }
      return opts.oldTargetId;
    };

    currentTargetId = pickReplacement(await opts.listTabs());
    if (currentTargetId === opts.oldTargetId) {
      await new Promise((r) => setTimeout(r, 800));
      currentTargetId = pickReplacement(await opts.listTabs(), {
        allowSingleTabFallback: true,
      });
    }
  } catch {
    // Best-effort: fall back to pre-navigation targetId
  }
  return currentTargetId;
}

export function registerBrowserAgentSnapshotRoutes(
  app: BrowserRouteRegistrar,
  ctx: BrowserRouteContext,
) {
  app.post("/navigate", async (req, res) => {
    const body = readBody(req);
    const url = toStringOrEmpty(body.url);
    const targetId = toStringOrEmpty(body.targetId) || undefined;
    if (!url) {
      return jsonError(res, 400, "url is required");
    }
    await withRouteTabContext({
      req,
      res,
      ctx,
      targetId,
      run: async ({ profileCtx, tab, cdpUrl }) => {
        if (getBrowserProfileCapabilities(profileCtx.profile).usesChromeMcp) {
          const ssrfPolicyOpts = withBrowserNavigationPolicy(ctx.state().resolved.ssrfPolicy);
          await assertBrowserNavigationAllowed({ url, ...ssrfPolicyOpts });
          const result = await navigateChromeMcpPage({
            profileName: profileCtx.profile.name,
            userDataDir: profileCtx.profile.userDataDir,
            targetId: tab.targetId,
            url,
          });
          await assertBrowserNavigationResultAllowed({ url: result.url, ...ssrfPolicyOpts });
          return res.json({ ok: true, targetId: tab.targetId, ...result });
        }
        const pw = await requirePwAi(res, "navigate");
        if (!pw) {
          return;
        }
        const result = await pw.navigateViaPlaywright({
          cdpUrl,
          targetId: tab.targetId,
          url,
          ...withBrowserNavigationPolicy(ctx.state().resolved.ssrfPolicy),
        });
        const currentTargetId = await resolveTargetIdAfterNavigate({
          oldTargetId: tab.targetId,
          navigatedUrl: result.url,
          listTabs: () => profileCtx.listTabs(),
        });
        res.json({ ok: true, targetId: currentTargetId, ...result });
      },
    });
  });

  app.post("/pdf", async (req, res) => {
    const body = readBody(req);
    const targetId = toStringOrEmpty(body.targetId) || undefined;
    const profileCtx = resolveProfileContext(req, res, ctx);
    if (!profileCtx) {
      return;
    }
    if (getBrowserProfileCapabilities(profileCtx.profile).usesChromeMcp) {
      return jsonError(
        res,
        501,
        "pdf is not supported for existing-session profiles yet; use screenshot/snapshot instead.",
      );
    }
    await withPlaywrightRouteContext({
      req,
      res,
      ctx,
      targetId,
      feature: "pdf",
      run: async ({ cdpUrl, tab, pw }) => {
        const pdf = await pw.pdfViaPlaywright({
          cdpUrl,
          targetId: tab.targetId,
        });
        await saveBrowserMediaResponse({
          res,
          buffer: pdf.buffer,
          contentType: "application/pdf",
          maxBytes: pdf.buffer.byteLength,
          targetId: tab.targetId,
          url: tab.url,
        });
      },
    });
  });

  app.post("/screenshot", async (req, res) => {
    const body = readBody(req);
    const targetId = toStringOrEmpty(body.targetId) || undefined;
    const fullPage = toBoolean(body.fullPage) ?? false;
    const ref = toStringOrEmpty(body.ref) || undefined;
    const element = toStringOrEmpty(body.element) || undefined;
    const type = body.type === "jpeg" ? "jpeg" : "png";

    if (fullPage && (ref || element)) {
      return jsonError(res, 400, "fullPage is not supported for element screenshots");
    }

    await withRouteTabContext({
      req,
      res,
      ctx,
      targetId,
      run: async ({ profileCtx, tab, cdpUrl }) => {
        if (getBrowserProfileCapabilities(profileCtx.profile).usesChromeMcp) {
          if (element) {
            return jsonError(
              res,
              400,
              "element screenshots are not supported for existing-session profiles; use ref from snapshot.",
            );
          }
          const buffer = await takeChromeMcpScreenshot({
            profileName: profileCtx.profile.name,
            userDataDir: profileCtx.profile.userDataDir,
            targetId: tab.targetId,
            uid: ref,
            fullPage,
            format: type,
          });
          await saveNormalizedScreenshotResponse({
            res,
            buffer,
            type,
            targetId: tab.targetId,
            url: tab.url,
          });
          return;
        }

        let buffer: Buffer;
        const shouldUsePlaywright = shouldUsePlaywrightForScreenshot({
          profile: profileCtx.profile,
          wsUrl: tab.wsUrl,
          ref,
          element,
        });
        if (shouldUsePlaywright) {
          const pw = await requirePwAi(res, "screenshot");
          if (!pw) {
            return;
          }
          const snap = await pw.takeScreenshotViaPlaywright({
            cdpUrl,
            targetId: tab.targetId,
            ref,
            element,
            fullPage,
            type,
          });
          buffer = snap.buffer;
        } else {
          buffer = await captureScreenshot({
            wsUrl: tab.wsUrl ?? "",
            fullPage,
            format: type,
            quality: type === "jpeg" ? 85 : undefined,
          });
        }

        await saveNormalizedScreenshotResponse({
          res,
          buffer,
          type,
          targetId: tab.targetId,
          url: tab.url,
        });
      },
    });
  });

  app.get("/snapshot", async (req, res) => {
    const profileCtx = resolveProfileContext(req, res, ctx);
    if (!profileCtx) {
      return;
    }
    const targetId = typeof req.query.targetId === "string" ? req.query.targetId.trim() : "";
    const hasPlaywright = Boolean(await getPwAiModule());
    const plan = resolveSnapshotPlan({
      profile: profileCtx.profile,
      query: req.query,
      hasPlaywright,
    });

    try {
      const tab = await profileCtx.ensureTabAvailable(targetId || undefined);
      if ((plan.labels || plan.mode === "efficient") && plan.format === "aria") {
        return jsonError(res, 400, "labels/mode=efficient require format=ai");
      }
      if (getBrowserProfileCapabilities(profileCtx.profile).usesChromeMcp) {
        if (plan.selectorValue || plan.frameSelectorValue) {
          return jsonError(
            res,
            400,
            "selector/frame snapshots are not supported for existing-session profiles; snapshot the whole page and use refs.",
          );
        }
        const snapshot = await takeChromeMcpSnapshot({
          profileName: profileCtx.profile.name,
          userDataDir: profileCtx.profile.userDataDir,
          targetId: tab.targetId,
        });
        if (plan.format === "aria") {
          return res.json({
            ok: true,
            format: "aria",
            targetId: tab.targetId,
            url: tab.url,
            nodes: flattenChromeMcpSnapshotToAriaNodes(snapshot, plan.limit),
          });
        }
        const built = buildAiSnapshotFromChromeMcpSnapshot({
          root: snapshot,
          options: {
            interactive: plan.interactive ?? undefined,
            compact: plan.compact ?? undefined,
            maxDepth: plan.depth ?? undefined,
          },
          maxChars: plan.resolvedMaxChars,
        });
        if (plan.labels) {
          const refs = Object.keys(built.refs);
          const labelResult = await renderChromeMcpLabels({
            profileName: profileCtx.profile.name,
            userDataDir: profileCtx.profile.userDataDir,
            targetId: tab.targetId,
            refs,
          });
          try {
            const labeled = await takeChromeMcpScreenshot({
              profileName: profileCtx.profile.name,
              userDataDir: profileCtx.profile.userDataDir,
              targetId: tab.targetId,
              format: "png",
            });
            const normalized = await normalizeBrowserScreenshot(labeled, {
              maxSide: DEFAULT_BROWSER_SCREENSHOT_MAX_SIDE,
              maxBytes: DEFAULT_BROWSER_SCREENSHOT_MAX_BYTES,
            });
            await ensureMediaDir();
            const saved = await saveMediaBuffer(
              normalized.buffer,
              normalized.contentType ?? "image/png",
              "browser",
              DEFAULT_BROWSER_SCREENSHOT_MAX_BYTES,
            );
            return res.json({
              ok: true,
              format: "ai",
              targetId: tab.targetId,
              url: tab.url,
              labels: true,
              labelsCount: labelResult.labels,
              labelsSkipped: labelResult.skipped,
              imagePath: path.resolve(saved.path),
              imageType: normalized.contentType?.includes("jpeg") ? "jpeg" : "png",
              ...built,
            });
          } finally {
            await clearChromeMcpOverlay({
              profileName: profileCtx.profile.name,
              userDataDir: profileCtx.profile.userDataDir,
              targetId: tab.targetId,
            });
          }
        }
        return res.json({
          ok: true,
          format: "ai",
          targetId: tab.targetId,
          url: tab.url,
          ...built,
        });
      }
      if (plan.format === "ai") {
        const pw = await requirePwAi(res, "ai snapshot");
        if (!pw) {
          return;
        }
        const roleSnapshotArgs = {
          cdpUrl: profileCtx.profile.cdpUrl,
          targetId: tab.targetId,
          selector: plan.selectorValue,
          frameSelector: plan.frameSelectorValue,
          refsMode: plan.refsMode,
          options: {
            interactive: plan.interactive ?? undefined,
            compact: plan.compact ?? undefined,
            maxDepth: plan.depth ?? undefined,
          },
        };

        const snap = plan.wantsRoleSnapshot
          ? await pw.snapshotRoleViaPlaywright(roleSnapshotArgs)
          : await pw
              .snapshotAiViaPlaywright({
                cdpUrl: profileCtx.profile.cdpUrl,
                targetId: tab.targetId,
                ...(typeof plan.resolvedMaxChars === "number"
                  ? { maxChars: plan.resolvedMaxChars }
                  : {}),
              })
              .catch(async (err) => {
                // Public-API fallback when Playwright's private _snapshotForAI is missing.
                if (String(err).toLowerCase().includes("_snapshotforai")) {
                  return await pw.snapshotRoleViaPlaywright(roleSnapshotArgs);
                }
                throw err;
              });
        if (plan.labels) {
          const labeled = await pw.screenshotWithLabelsViaPlaywright({
            cdpUrl: profileCtx.profile.cdpUrl,
            targetId: tab.targetId,
            refs: "refs" in snap ? snap.refs : {},
            type: "png",
          });
          const normalized = await normalizeBrowserScreenshot(labeled.buffer, {
            maxSide: DEFAULT_BROWSER_SCREENSHOT_MAX_SIDE,
            maxBytes: DEFAULT_BROWSER_SCREENSHOT_MAX_BYTES,
          });
          await ensureMediaDir();
          const saved = await saveMediaBuffer(
            normalized.buffer,
            normalized.contentType ?? "image/png",
            "browser",
            DEFAULT_BROWSER_SCREENSHOT_MAX_BYTES,
          );
          const imageType = normalized.contentType?.includes("jpeg") ? "jpeg" : "png";
          return res.json({
            ok: true,
            format: plan.format,
            targetId: tab.targetId,
            url: tab.url,
            labels: true,
            labelsCount: labeled.labels,
            labelsSkipped: labeled.skipped,
            imagePath: path.resolve(saved.path),
            imageType,
            ...snap,
          });
        }

        return res.json({
          ok: true,
          format: plan.format,
          targetId: tab.targetId,
          url: tab.url,
          ...snap,
        });
      }

      const snap = shouldUsePlaywrightForAriaSnapshot({
        profile: profileCtx.profile,
        wsUrl: tab.wsUrl,
      })
        ? (() => {
            // Extension relay doesn't expose per-page WS URLs; run AX snapshot via Playwright CDP session.
            // Also covers cases where wsUrl is missing/unusable.
            return requirePwAi(res, "aria snapshot").then(async (pw) => {
              if (!pw) {
                return null;
              }
              return await pw.snapshotAriaViaPlaywright({
                cdpUrl: profileCtx.profile.cdpUrl,
                targetId: tab.targetId,
                limit: plan.limit,
              });
            });
          })()
        : snapshotAria({ wsUrl: tab.wsUrl ?? "", limit: plan.limit });

      const resolved = await Promise.resolve(snap);
      if (!resolved) {
        return;
      }
      return res.json({
        ok: true,
        format: plan.format,
        targetId: tab.targetId,
        url: tab.url,
        ...resolved,
      });
    } catch (err) {
      handleRouteError(ctx, res, err);
    }
  });
}