import { expect, test, type Page } from "@playwright/test";
import { installRoutes } from "./helpers";

// The integrated grounding workspace: pick the AI target (whole search ⇄ a
// hand-picked selection), the [Search|Chat|Split] view modes, saved chats, and a
// pausable/resumable sweep. Retrieval runs client-side over the mocked fixtures;
// the AI provider is mocked here (native tool-calling), reusing the same SSE/
// tool_use shapes as ask-chat.spec.

type SLeaf = { k: "l"; q: string; s: "transcripts" | "chat" | "metadata" };
type SGroup = { k: "g"; o: "AND" | "OR"; c: (SLeaf | SGroup)[] };
const qt = (root: SGroup) => encodeURIComponent(JSON.stringify(root));

// transcripts:"alpha" matches all three fixtures.
const ALL_TREE: SGroup = {
  k: "g",
  o: "AND",
  c: [{ k: "l", q: "alpha", s: "transcripts" }],
};

const CORS = {
  "access-control-allow-origin": "*",
  "access-control-allow-headers": "*",
  "access-control-allow-methods": "*",
};

function sse(text: string): string {
  return (
    `data: ${JSON.stringify({ type: "content_block_delta", delta: { type: "text_delta", text } })}\n\n` +
    `data: ${JSON.stringify({ type: "message_stop" })}\n\n`
  );
}
function toolUse(name: string, input: Record<string, unknown>) {
  return JSON.stringify({
    content: [{ type: "tool_use", id: "tu_1", name, input }],
    stop_reason: "tool_use",
  });
}

// A normal chat mock: native gather searches once then finishes; answers stream
// Markdown, echoing which grounded titles reached the answer request so a test can
// assert WHAT the turn was grounded in.
async function mockChat(page: Page, answerBodies: string[] = []) {
  await page.route("https://api.anthropic.com/**", async (route) => {
    if (route.request().method() === "OPTIONS") {
      await route.fulfill({ status: 204, headers: CORS });
      return;
    }
    const body = route.request().postDataJSON() as {
      system?: string;
      tools?: unknown[];
      messages?: { role: string; content: unknown }[];
    };
    const system = body.system ?? "";
    if (system.includes("Markdown")) {
      answerBodies.push(JSON.stringify(body.messages ?? []));
      await route.fulfill({
        status: 200,
        headers: { ...CORS, "content-type": "text/event-stream" },
        body: sse("Grounded answer [1]"),
      });
      return;
    }
    if (Array.isArray(body.tools)) {
      const msgs = body.messages ?? [];
      const lastUser = [...msgs].reverse().find((m) => m.role === "user");
      const isToolResult = Array.isArray(lastUser?.content);
      await route.fulfill({
        status: 200,
        headers: { ...CORS, "content-type": "application/json" },
        body: isToolResult ? toolUse("finish", {}) : toolUse("finish", {}),
      });
      return;
    }
    await route.fulfill({
      status: 200,
      headers: { ...CORS, "content-type": "text/event-stream" },
      body: sse("DONE"),
    });
  });
}

async function keyIn(page: Page, mode: "Scripted" | "Native tools" = "Native tools") {
  const key = page.locator('input[placeholder^="sk-ant"]');
  // Clear first: a remembered key may already be restored (e.g. after a reload),
  // and a bare fill() would append to it.
  await key.clear();
  await key.fill("sk-ant-test");
  await expect(key).toHaveValue("sk-ant-test");
  await page.getByRole("button", { name: mode }).click();
}

// Run a real search on `/`, wait for the cards, cross to the chat.
async function searchThenChat(page: Page, cards: number) {
  await page.goto(`/?qt=${qt(ALL_TREE)}`);
  await expect(page.locator("[data-card-header]")).toHaveCount(cards, {
    timeout: 15_000,
  });
  await page.getByTestId("workspace-nav").getByRole("link", { name: "Chat" }).click();
}

test.describe("grounding workspace", () => {
  test("selecting a subset grounds the chat in exactly those videos", async ({
    page,
  }) => {
    await installRoutes(page);
    const answerBodies: string[] = [];
    await mockChat(page, answerBodies);
    await page.goto(`/?qt=${qt(ALL_TREE)}`);
    await expect(page.locator("[data-card-header]")).toHaveCount(3, { timeout: 15_000 });

    // Select 2 of the 3 results via the per-card checkboxes.
    await page.getByRole("checkbox", { name: /Transcript only/ }).check();
    await page.getByRole("checkbox", { name: /Large live chat/ }).check();
    await expect(page.getByTestId("selection-toolbar")).toContainText("2 selected");

    // Cross to the chat — it auto-grounds; flip the palette to Selection.
    await page.getByTestId("workspace-nav").getByRole("link", { name: "Chat" }).click();
    await keyIn(page);
    await page
      .getByTestId("grounding-mode")
      .getByRole("tab", { name: /Selection/ })
      .click();

    // The palette + panel report the 2-video selection, shown as chips.
    await expect(page.getByTestId("grounding-palette")).toContainText(
      "Grounded in 2 selected videos",
    );
    const chips = page.getByTestId("selection-chips");
    await expect(chips).toContainText("Transcript only");
    await expect(chips).toContainText("Large live chat");
    await expect(chips).not.toContainText("Small live chat");

    // A strict ask grounds ONLY in the 2 selected — the answer request carries
    // their titles but not the deselected "Small live chat".
    await page.getByPlaceholder(/Ask about the transcripts/).fill("summarize these");
    await page.getByRole("button", { name: "Ask", exact: true }).click();
    await expect(page.getByText(/Grounded answer/)).toBeVisible();
    const last = answerBodies[answerBodies.length - 1];
    expect(last).toContain("Transcript only");
    expect(last).toContain("Large live chat");
    expect(last).not.toContain("Small live chat");
  });

  test("a per-card Ask grounds in that one video", async ({ page }) => {
    await installRoutes(page);
    await mockChat(page);
    await page.goto(`/?qt=${qt(ALL_TREE)}`);
    await expect(page.locator("[data-card-header]")).toHaveCount(3, { timeout: 15_000 });

    // Click "Ask" on the Small live chat card → navigates to the chat, grounded
    // in exactly that one video.
    await page
      .locator("[data-card-header]", { hasText: "Small live chat" })
      .getByRole("button", { name: "Ask" })
      .click();
    await expect(page).toHaveURL(/\/ask\/?$/);
    await keyIn(page);
    await expect(page.getByTestId("grounding-palette")).toContainText(
      "Grounded in 1 selected video",
    );
    await expect(page.getByTestId("selection-chips")).toContainText("Small live chat");
  });

  test("the [Search|Chat|Split] control fullscreens either pane and degrades to tabs when narrow", async ({
    page,
  }) => {
    await installRoutes(page);
    await mockChat(page);
    await page.goto(`/?qt=${qt(ALL_TREE)}`);
    await expect(page.locator("[data-card-header]").first()).toBeVisible({ timeout: 15_000 });

    const nav = page.getByTestId("workspace-nav");
    const searchPane = page.getByTestId("pane-search");
    const chatPane = page.getByTestId("pane-chat");

    // Search view (route `/`): results visible, chat hidden.
    await expect(searchPane).toBeVisible();
    await expect(chatPane).toBeHidden();

    // Split (wide viewport) → BOTH panes visible side by side.
    await nav.getByTestId("view-split").click();
    await expect(searchPane).toBeVisible();
    await expect(chatPane).toBeVisible();

    // Chat → chat pane fullscreen, search hidden.
    await nav.getByRole("link", { name: "Chat" }).click();
    await expect(chatPane).toBeVisible();
    await expect(searchPane).toBeHidden();

    // Narrow viewport → Split isn't offered (degrades to Search ⇄ Chat tabs).
    await page.setViewportSize({ width: 600, height: 900 });
    await expect(nav.getByTestId("view-split")).toHaveCount(0);
    await expect(nav.getByRole("link", { name: "Search" })).toBeVisible();
    await expect(nav.getByRole("link", { name: "Chat" })).toBeVisible();
  });

  test("save a chat, New chat, reload, then load it → the conversation is restored", async ({
    page,
  }) => {
    await installRoutes(page);
    await mockChat(page);
    await searchThenChat(page, 3);
    await keyIn(page);

    // Produce an answer, then Save as… (browser prompt → "research").
    await page.getByPlaceholder(/Ask about the transcripts/).fill("what happened");
    await page.getByRole("button", { name: "Ask", exact: true }).click();
    await expect(page.getByText(/Grounded answer/)).toBeVisible();

    page.once("dialog", (d) => d.accept("research"));
    await page.getByTestId("saved-chats-row").getByRole("button", { name: "Save as…" }).click();
    await expect(page.getByTestId("saved-chat-select")).toHaveValue("research");

    // New chat clears the working conversation.
    await page.getByRole("button", { name: "New chat" }).click();
    await expect(page.getByText(/Grounded answer/)).toHaveCount(0);

    // Reload — the saved chat is still in the library; load it → restored.
    await page.waitForTimeout(400);
    await page.reload();
    await keyIn(page);
    await page
      .getByTestId("saved-chat-select")
      .selectOption("research");
    await expect(page.getByText(/Grounded answer/)).toBeVisible();
    await expect(page.getByText("what happened")).toBeVisible();
  });

  // A transient 429 (with a Retry-After hint) is now RETRIED below the sweep
  // driver — the run paces itself and carries on, no pause banner. Only a 429 that
  // outlasts the retry cap falls through to the pause/checkpoint (next test).
  test("a transient 429 (Retry-After) mid-sweep auto-recovers and completes, with no pause", async ({
    page,
  }) => {
    await installRoutes(page);
    // Force batch size 2 so the 3 fixtures span 2 batches.
    await page.addInitScript(() => localStorage.setItem("ytdlp-tb:ai:sweepchunk", "2"));

    let batch2FirstRoundHits = 0;
    await page.route("https://api.anthropic.com/**", async (route) => {
      if (route.request().method() === "OPTIONS") {
        await route.fulfill({ status: 204, headers: CORS });
        return;
      }
      const body = route.request().postDataJSON() as {
        system?: string;
        tools?: unknown[];
        messages?: { role: string; content: unknown }[];
      };
      const system = body.system ?? "";
      const msgs = body.messages ?? [];
      const lastUser = [...msgs].reverse().find((m) => m.role === "user");
      const isToolResult = Array.isArray(lastUser?.content);
      const text = typeof lastUser?.content === "string" ? lastUser.content : "";
      if (Array.isArray(body.tools) && /running report/.test(system)) {
        const batch = /Batch (\d+) of/.exec(text)?.[1] ?? "1";
        // Batch 2's first request 429s ONCE, with a 1s Retry-After hint. The
        // limiter waits the hint and re-issues; the retry succeeds.
        if (batch === "2" && !isToolResult) {
          batch2FirstRoundHits += 1;
          if (batch2FirstRoundHits === 1) {
            await route.fulfill({
              status: 429,
              headers: { ...CORS, "content-type": "application/json", "retry-after": "1" },
              body: JSON.stringify({ type: "error", error: { message: "rate limited" } }),
            });
            return;
          }
        }
        if (isToolResult) {
          await route.fulfill({
            status: 200,
            headers: { ...CORS, "content-type": "application/json" },
            body: toolUse("finish", {}),
          });
          return;
        }
        await route.fulfill({
          status: 200,
          headers: { ...CORS, "content-type": "application/json" },
          body: toolUse("update_report", {
            section: `Batch ${batch} finding`,
            content: `Finding recorded in batch ${batch}.`,
          }),
        });
        return;
      }
      await route.fulfill({
        status: 200,
        headers: { ...CORS, "content-type": "text/event-stream" },
        body: sse("DONE"),
      });
    });

    await searchThenChat(page, 3);
    await keyIn(page);
    await expect(page.getByText(/Grounded in 3 results/)).toBeVisible();

    await page.getByRole("button", { name: /Build report from all 3 results/ }).click();

    // Both batches fold in — the 429 was absorbed by a retry, not a pause.
    await expect(page.getByText("Finding recorded in batch 1.")).toBeVisible({ timeout: 15_000 });
    await expect(page.getByText("Finding recorded in batch 2.")).toBeVisible({ timeout: 15_000 });
    // No pause banner, and the sweep reported completion.
    await expect(page.getByTestId("sweep-paused")).toHaveCount(0);
    await expect(page.getByText(/Built a report from 3 results/)).toBeVisible();
  });

  // A 429 that persists past the retry cap is the genuine last resort: it falls
  // through to the existing pause/checkpoint so the user can Resume later (e.g. a
  // daily-quota cap that clears tomorrow). Guards that regression path.
  test("a persistent 429 past the retry cap PAUSES; Resume — even after a model switch — finishes it", async ({
    page,
  }) => {
    await installRoutes(page);
    await page.addInitScript(() => localStorage.setItem("ytdlp-tb:ai:sweepchunk", "2"));

    let batch2Hits = 0;
    await page.route("https://api.anthropic.com/**", async (route) => {
      if (route.request().method() === "OPTIONS") {
        await route.fulfill({ status: 204, headers: CORS });
        return;
      }
      const body = route.request().postDataJSON() as {
        system?: string;
        tools?: unknown[];
        messages?: { role: string; content: unknown }[];
      };
      const system = body.system ?? "";
      const msgs = body.messages ?? [];
      const lastUser = [...msgs].reverse().find((m) => m.role === "user");
      const isToolResult = Array.isArray(lastUser?.content);
      const text = typeof lastUser?.content === "string" ? lastUser.content : "";
      if (Array.isArray(body.tools) && /running report/.test(system)) {
        const batch = /Batch (\d+) of/.exec(text)?.[1] ?? "1";
        // Batch 2 429s on every attempt of the INITIAL run (enough to exhaust the
        // retry cap → pause). The later resumed run's attempts succeed.
        if (batch === "2" && !isToolResult) {
          batch2Hits += 1;
          if (batch2Hits <= 4) {
            await route.fulfill({
              status: 429,
              headers: { ...CORS, "content-type": "application/json", "retry-after": "1" },
              body: JSON.stringify({ type: "error", error: { message: "rate limited" } }),
            });
            return;
          }
        }
        if (isToolResult) {
          await route.fulfill({
            status: 200,
            headers: { ...CORS, "content-type": "application/json" },
            body: toolUse("finish", {}),
          });
          return;
        }
        await route.fulfill({
          status: 200,
          headers: { ...CORS, "content-type": "application/json" },
          body: toolUse("update_report", {
            section: `Batch ${batch} finding`,
            content: `Finding recorded in batch ${batch}.`,
          }),
        });
        return;
      }
      await route.fulfill({
        status: 200,
        headers: { ...CORS, "content-type": "text/event-stream" },
        body: sse("DONE"),
      });
    });

    await searchThenChat(page, 3);
    await keyIn(page);
    await expect(page.getByText(/Grounded in 3 results/)).toBeVisible();

    await page.getByRole("button", { name: /Build report from all 3 results/ }).click();

    // Batch 1 folds in; batch 2's retries all 429 → the run pauses (not errors).
    await expect(page.getByText("Finding recorded in batch 1.")).toBeVisible({ timeout: 15_000 });
    const paused = page.getByTestId("sweep-paused");
    await expect(paused).toBeVisible({ timeout: 15_000 });
    await expect(paused).toContainText("usage limit");
    // The partial report survived.
    await expect(page.getByText("Finding recorded in batch 2.")).toHaveCount(0);

    // Switch the model (the checkpoint carries no model binding) and Resume.
    const modelInput = page.locator('input[list="ask-models"]');
    await modelInput.fill("claude-sonnet-5");
    await paused.getByRole("button", { name: "Resume" }).click();

    // The resumed batch completes and the report is now whole.
    await expect(page.getByText("Finding recorded in batch 2.")).toBeVisible({ timeout: 15_000 });
    await expect(page.getByTestId("sweep-paused")).toHaveCount(0);
    await expect(page.getByText("Finding recorded in batch 1.")).toBeVisible();
  });
});