import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import type { DiagramBundle } from "../services/renderers/types.js";

// ─── Mock playwright-core chromium ─────────────────────────────────────────────

const mockPageSetContent = vi.fn().mockResolvedValue(undefined);
const mockPageWaitForSelector = vi.fn().mockResolvedValue(undefined);
const mockPageEval = vi.fn().mockResolvedValue(
  '<svg xmlns="http://www.w3.org/2000/svg"><rect width="100" height="50"/></svg>',
);
const mockBrowserClose = vi.fn().mockResolvedValue(undefined);
const mockPage = {
  setContent: mockPageSetContent,
  waitForSelector: mockPageWaitForSelector,
  $eval: mockPageEval,
};
const mockBrowser = {
  newPage: vi.fn().mockResolvedValue(mockPage),
  close: mockBrowserClose,
};
const mockChromiumLaunch = vi.fn().mockResolvedValue(mockBrowser);

vi.mock("playwright-core", () => ({
  chromium: {
    launch: mockChromiumLaunch,
  },
}));

// ─── Mock LLM inference ─────────────────────────────────────────────────────────

const MOCK_MERMAID_SOURCE = "graph TD\n  A[Login]-->B[Validate]\n  B-->C[Dashboard]";
const mockPuterChatComplete = vi.fn().mockResolvedValue(MOCK_MERMAID_SOURCE);

vi.mock("../services/puter-inference.js", () => ({
  puterChatComplete: (...args: unknown[]) => mockPuterChatComplete(...args),
}));

// ─── Tests ──────────────────────────────────────────────────────────────────────

describe("stripUnsafeDirectives", () => {
  let stripUnsafeDirectives: (source: string) => { cleaned: string; stripped: boolean };

  beforeEach(async () => {
    // Re-import to get a fresh module each test
    const mod = await import("../services/renderers/diagram-renderer.js");
    stripUnsafeDirectives = mod.stripUnsafeDirectives;
  });

  it("strips %%{init}%% blocks and marks stripped=true", () => {
    const input = '%%{init: {"theme": "dark"}}%%\ngraph TD\n  A-->B';
    const result = stripUnsafeDirectives(input);
    expect(result.stripped).toBe(true);
    expect(result.cleaned).not.toContain("%%{");
    expect(result.cleaned).toContain("graph TD");
  });

  it('strips click "url" lines and marks stripped=true', () => {
    const input = 'graph TD\n  A-->B\n  click A "https://evil.com"';
    const result = stripUnsafeDirectives(input);
    expect(result.stripped).toBe(true);
    expect(result.cleaned).not.toContain("click A");
    expect(result.cleaned).toContain("graph TD");
  });

  it("strips click call fn() lines and marks stripped=true", () => {
    const input = "graph TD\n  A-->B\n  click A call myFn()";
    const result = stripUnsafeDirectives(input);
    expect(result.stripped).toBe(true);
    expect(result.cleaned).not.toContain("click A");
  });

  it("leaves clean source unchanged with stripped=false", () => {
    const input = "graph TD\n  A-->B";
    const result = stripUnsafeDirectives(input);
    expect(result.stripped).toBe(false);
    expect(result.cleaned).toBe("graph TD\n  A-->B");
  });

  it("strips both init and click directives simultaneously", () => {
    const input = '%%{init: {"theme": "dark"}}%%\ngraph TD\n  A-->B\n  click A "https://evil.com"';
    const result = stripUnsafeDirectives(input);
    expect(result.stripped).toBe(true);
    expect(result.cleaned).not.toContain("%%{");
    expect(result.cleaned).not.toContain("click A");
    expect(result.cleaned).toContain("graph TD");
  });
});

describe("buildDiagramPrompt", () => {
  let buildDiagramPrompt: (
    description: string,
    diagramType: string,
  ) => { system: string; user: string };

  beforeEach(async () => {
    const mod = await import("../services/renderers/diagram-renderer.js");
    buildDiagramPrompt = mod.buildDiagramPrompt;
  });

  it("includes 'flowchart' when diagramType is flowchart", () => {
    const result = buildDiagramPrompt("a login flow", "flowchart");
    const combined = result.system + result.user;
    expect(combined.toLowerCase()).toContain("flowchart");
  });

  it("includes 'architecture' when diagramType is architecture", () => {
    const result = buildDiagramPrompt("a microservices setup", "architecture");
    const combined = result.system + result.user;
    expect(combined.toLowerCase()).toContain("architecture");
  });

  it("includes the user's natural language description in the user prompt", () => {
    const desc = "unique description for test xyz123";
    const result = buildDiagramPrompt(desc, "flowchart");
    expect(result.user).toContain(desc);
  });

  it("instructs LLM to output ONLY valid Mermaid syntax without markdown fences", () => {
    const result = buildDiagramPrompt("some flow", "flowchart");
    const system = result.system.toLowerCase();
    expect(system).toContain("mermaid");
    expect(system).toMatch(/only|no markdown|no explanation/i);
  });

  it("includes sequence-specific preamble for sequence diagrams", () => {
    const result = buildDiagramPrompt("an API request flow", "sequence");
    const combined = result.system + result.user;
    expect(combined.toLowerCase()).toContain("sequence");
  });

  it("includes ERD-specific preamble for erd diagrams", () => {
    const result = buildDiagramPrompt("user-post relationship", "erd");
    const combined = result.system + result.user;
    expect(combined.toLowerCase()).toContain("erd");
  });

  it("includes mindmap-specific preamble for mindmap diagrams", () => {
    const result = buildDiagramPrompt("project ideas", "mindmap");
    const combined = result.system + result.user;
    expect(combined.toLowerCase()).toContain("mindmap");
  });
});

describe("renderDiagram integration", () => {
  let renderDiagram: (input: Record<string, unknown>) => Promise<import("../services/renderers/types.js").RenderResult>;

  beforeEach(async () => {
    vi.clearAllMocks();
    // Restore default mocks
    mockPuterChatComplete.mockResolvedValue(MOCK_MERMAID_SOURCE);
    mockChromiumLaunch.mockResolvedValue(mockBrowser);
    mockBrowser.newPage.mockResolvedValue(mockPage);
    mockPageEval.mockResolvedValue(
      '<svg xmlns="http://www.w3.org/2000/svg"><rect width="100" height="50"/></svg>',
    );
    mockBrowserClose.mockResolvedValue(undefined);

    const mod = await import("../services/renderers/diagram-renderer.js");
    renderDiagram = mod.renderDiagram;
  });

  afterEach(() => {
    vi.clearAllMocks();
  });

  it("returns a RenderResult with diagram-bundle JSON structure", async () => {
    const result = await renderDiagram({
      prompt: "A login flow with validation",
      diagramType: "flowchart",
      darkMode: false,
    });

    expect(result.filename).toBe("diagram-bundle.json");
    expect(result.contentType).toBe("application/json");

    const bundle = JSON.parse(result.buffer.toString()) as DiagramBundle;
    expect(bundle.type).toBe("diagram-bundle");
    expect(typeof bundle.svgBase64).toBe("string");
    expect(typeof bundle.pngBase64).toBe("string");
    expect(typeof bundle.mermaidSource).toBe("string");
    expect(typeof bundle.stripped).toBe("boolean");
  });

  it("mermaidSource in bundle is the LLM-generated Mermaid, not the original prompt", async () => {
    const result = await renderDiagram({
      prompt: "A login flow with validation",
      diagramType: "flowchart",
      darkMode: false,
    });
    const bundle = JSON.parse(result.buffer.toString()) as DiagramBundle;
    // The mermaidSource should contain the LLM-returned Mermaid, not the English prompt
    expect(bundle.mermaidSource).toContain("graph TD");
    expect(bundle.mermaidSource).not.toBe("A login flow with validation");
  });

  it("calls browser.close() for cleanup", async () => {
    await renderDiagram({
      prompt: "test prompt",
      diagramType: "flowchart",
      darkMode: false,
    });
    expect(mockBrowserClose).toHaveBeenCalled();
  });

  it("calls browser.close() even when rendering throws", async () => {
    mockPageWaitForSelector.mockRejectedValueOnce(new Error("timeout"));
    await expect(
      renderDiagram({ prompt: "bad diagram", diagramType: "flowchart", darkMode: false }),
    ).rejects.toThrow();
    expect(mockBrowserClose).toHaveBeenCalled();
  });

  it("uses source directly when provided (re-render path, skips LLM)", async () => {
    mockPuterChatComplete.mockClear();
    const result = await renderDiagram({
      source: "graph TD\n  X-->Y",
      diagramType: "flowchart",
      darkMode: false,
    });
    // LLM should NOT have been called
    expect(mockPuterChatComplete).not.toHaveBeenCalled();
    const bundle = JSON.parse(result.buffer.toString()) as DiagramBundle;
    expect(bundle.mermaidSource).toContain("graph TD");
  });

  it("sets stripped=true in bundle when LLM-generated source contains click directives", async () => {
    mockPuterChatComplete.mockResolvedValueOnce(
      'graph TD\n  A-->B\n  click A "https://evil.com"',
    );
    const result = await renderDiagram({
      prompt: "diagram with click",
      diagramType: "flowchart",
      darkMode: false,
    });
    const bundle = JSON.parse(result.buffer.toString()) as DiagramBundle;
    expect(bundle.stripped).toBe(true);
  });
});
