import { describe, expect, it } from "vitest";

import { buildCopilotSystemPrompt, buildSelectedContext } from "./system-prompt";

import type { DiagramType } from "./types";

describe("buildCopilotSystemPrompt", () => {
  it("should include the diagram type in uppercase", () => {
    const prompt = buildCopilotSystemPrompt("bpmn");
    expect(prompt).toContain("BPMN");
    expect(prompt).toContain("Type: BPMN");
  });

  it("should include diagram description for each type", () => {
    const types: DiagramType[] = [
      "bpmn",
      "er",
      "orgchart",
      "architecture",
      "sequence",
      "flowchart",
    ];

    for (const type of types) {
      const prompt = buildCopilotSystemPrompt(type);
      expect(prompt).toContain(`Type: ${type.toUpperCase()}`);
      expect(prompt).toContain("domaingraph AI copilot");
    }
  });

  it("should include generation instructions with generateDiagram tool", () => {
    const prompt = buildCopilotSystemPrompt("bpmn");
    expect(prompt).toContain("generateDiagram");
    expect(prompt).toContain("Generate complete diagrams");
    expect(prompt).not.toContain("CHAT-ONLY");
    expect(prompt).not.toContain("cannot modify the diagram directly");
  });

  it("should include today's date", () => {
    const prompt = buildCopilotSystemPrompt("flowchart");
    const year = new Date().getFullYear().toString();
    expect(prompt).toContain(year);
  });

  it("should include ER-specific description", () => {
    const prompt = buildCopilotSystemPrompt("er");
    expect(prompt).toContain("Entity-Relationship");
    expect(prompt).toContain("entities");
  });

  it("should include architecture-specific description", () => {
    const prompt = buildCopilotSystemPrompt("architecture");
    expect(prompt).toContain("services");
    expect(prompt).toContain("databases");
  });

  it("should include BPMN-specific node types", () => {
    const prompt = buildCopilotSystemPrompt("bpmn");
    expect(prompt).toContain("activity");
    expect(prompt).toContain("gateway-exclusive");
    expect(prompt).toContain("gateway-parallel");
    expect(prompt).toContain("gateway-inclusive");
    expect(prompt).toContain("start-event");
    expect(prompt).toContain("end-event");
    expect(prompt).toContain("event-timer");
    expect(prompt).toContain("event-message");
    expect(prompt).toContain("subprocess");
    expect(prompt).toContain("data-object");
    expect(prompt).toContain("annotation");
  });

  it("should include BPMN-specific edge types", () => {
    const prompt = buildCopilotSystemPrompt("bpmn");
    expect(prompt).toContain("sequence");
    expect(prompt).toContain("message");
    expect(prompt).toContain("association");
  });

  it("should include BPMN pools and lanes guidance", () => {
    const prompt = buildCopilotSystemPrompt("bpmn");
    expect(prompt).toContain("pools");
    expect(prompt).toContain("lanes");
  });

  it("should include node types for each diagram type", () => {
    const typeNodeMap: Record<DiagramType, string[]> = {
      bpmn: ["activity", "gateway-exclusive", "start-event", "end-event"],
      er: ["entity"],
      orgchart: ["person"],
      architecture: ["service", "database", "queue", "loadbalancer", "external"],
      sequence: ["participant", "fragment"],
      flowchart: ["process", "decision", "terminal", "io"],
    };

    for (const [type, nodes] of Object.entries(typeNodeMap)) {
      const prompt = buildCopilotSystemPrompt(type as DiagramType);
      for (const node of nodes) {
        expect(prompt).toContain(node);
      }
    }
  });

  it("should include sequence-specific edge types", () => {
    const prompt = buildCopilotSystemPrompt("sequence");
    expect(prompt).toContain("sync");
    expect(prompt).toContain("async");
    expect(prompt).toContain("return");
  });

  it("should include ER column format with key fields", () => {
    const prompt = buildCopilotSystemPrompt("er");
    expect(prompt).toContain("columns");
    expect(prompt).toContain("isPrimaryKey");
    expect(prompt).toContain("isForeignKey");
    expect(prompt).toContain("isNullable");
  });

  it("should include type inference instructions for all types", () => {
    const prompt = buildCopilotSystemPrompt("bpmn");
    expect(prompt).toContain("infer");
    expect(prompt).toContain("bpmn");
    expect(prompt).toContain("er");
    expect(prompt).toContain("orgchart");
    expect(prompt).toContain("architecture");
    expect(prompt).toContain("sequence");
    expect(prompt).toContain("flowchart");
  });

  it("should include 200-node generation cap", () => {
    const prompt = buildCopilotSystemPrompt("bpmn");
    expect(prompt).toContain("200");
  });

  it("should include graph context when provided", () => {
    const graphContext = '{"nodes":[],"edges":[]}';
    const prompt = buildCopilotSystemPrompt("bpmn", { graphContext });
    expect(prompt).toContain(graphContext);
  });

  it("should show empty canvas when no graph context", () => {
    const prompt = buildCopilotSystemPrompt("bpmn");
    expect(prompt).toContain("Empty canvas");
  });

  it("should include a JSON example for each diagram type", () => {
    const types: DiagramType[] = [
      "bpmn",
      "er",
      "orgchart",
      "architecture",
      "sequence",
      "flowchart",
    ];

    for (const type of types) {
      const prompt = buildCopilotSystemPrompt(type);
      expect(prompt).toContain("```json");
      expect(prompt).toContain(`"diagramType":"${type}"`);
    }
  });

  it("should include GraphData format specification", () => {
    const prompt = buildCopilotSystemPrompt("bpmn");
    expect(prompt).toContain("meta");
    expect(prompt).toContain("nodes");
    expect(prompt).toContain("edges");
    expect(prompt).toContain("diagramType");
    expect(prompt).toContain("layoutDirection");
  });

  it("should include ID generation convention", () => {
    const prompt = buildCopilotSystemPrompt("flowchart");
    expect(prompt).toContain('"n1"');
    expect(prompt).toContain('"e1"');
  });

  describe("scoped context — selectedElements", () => {
    const selectedElements = [
      { id: "n1", type: "activity", label: "Validate order" },
      { id: "n2", type: "gateway-exclusive", label: "Valid?" },
    ];

    it("should include scoped context section when selectedElements provided", () => {
      const prompt = buildCopilotSystemPrompt("bpmn", { selectedElements });
      expect(prompt).toContain("Scoped context");
      expect(prompt).toContain("targeted modification");
    });

    it("should list all selected elements by label and type", () => {
      const prompt = buildCopilotSystemPrompt("bpmn", { selectedElements });
      expect(prompt).toContain("**Validate order** (activity)");
      expect(prompt).toContain("**Valid?** (gateway-exclusive)");
    });

    it("should show element count", () => {
      const prompt = buildCopilotSystemPrompt("bpmn", { selectedElements });
      expect(prompt).toContain("2 element(s)");
    });

    it("should include selectedContext JSON when provided", () => {
      const selectedContext = JSON.stringify({
        selectedNodes: selectedElements,
        connectedEdges: [{ id: "e1", from: "n1", to: "n2" }],
      });
      const prompt = buildCopilotSystemPrompt("bpmn", {
        selectedElements,
        selectedContext,
      });
      expect(prompt).toContain("Selected element details");
      expect(prompt).toContain(selectedContext);
    });

    it("should include preservation instructions for scoped mode", () => {
      const prompt = buildCopilotSystemPrompt("bpmn", { selectedElements });
      expect(prompt).toContain("Include ALL nodes and edges");
      expect(prompt).toContain("Preserve all other nodes and edges unchanged");
      expect(prompt).toContain("minimal changes");
    });

    it("should NOT include scoped context when selectedElements is empty", () => {
      const prompt = buildCopilotSystemPrompt("bpmn", { selectedElements: [] });
      expect(prompt).not.toContain("Scoped context");
    });

    it("should NOT include scoped context when selectedElements is undefined", () => {
      const prompt = buildCopilotSystemPrompt("bpmn", {});
      expect(prompt).not.toContain("Scoped context");
    });

    it("should include scoped context alongside graph context", () => {
      const graphContext = '{"nodes":[{"id":"n1"}],"edges":[]}';
      const prompt = buildCopilotSystemPrompt("bpmn", {
        graphContext,
        selectedElements,
      });
      expect(prompt).toContain(graphContext);
      expect(prompt).toContain("Scoped context");
    });
  });
});

interface ContextResult {
  selectedNodes: Array<{ id: string; [key: string]: unknown }>;
  connectedEdges?: Array<{ id: string; from: string; to: string }>;
  neighborNodes?: Array<{ id: string; label: string }>;
}

const parseContext = (raw: string): ContextResult => JSON.parse(raw) as ContextResult;

describe("buildSelectedContext", () => {
  const selectedElements = [
    { id: "n1", type: "activity", label: "Validate order" },
    { id: "n2", type: "gateway-exclusive", label: "Valid?" },
  ];

  it("should return selectedNodes-only JSON when no graphContext", () => {
    const result = parseContext(buildSelectedContext(selectedElements));
    expect(result.selectedNodes).toEqual(selectedElements);
    expect(result.connectedEdges).toBeUndefined();
    expect(result.neighborNodes).toBeUndefined();
  });

  it("should extract connected edges from graphContext", () => {
    const graphContext = JSON.stringify({
      nodes: [
        { id: "n1", label: "Validate order" },
        { id: "n2", label: "Valid?" },
        { id: "n3", label: "Reject" },
      ],
      edges: [
        { id: "e1", from: "n1", to: "n2" },
        { id: "e2", from: "n2", to: "n3" },
        { id: "e3", from: "n3", to: "n3" }, // self-loop on non-selected, not connected
      ],
    });
    const result = parseContext(buildSelectedContext(selectedElements, graphContext));
    expect(result.connectedEdges).toHaveLength(2);
    expect(result.connectedEdges!.map((e) => e.id)).toEqual(["e1", "e2"]);
  });

  it("should extract 1-hop neighbor nodes", () => {
    const graphContext = JSON.stringify({
      nodes: [
        { id: "n1", label: "Validate order" },
        { id: "n2", label: "Valid?" },
        { id: "n3", label: "Reject" },
        { id: "n4", label: "Unconnected" },
      ],
      edges: [{ id: "e1", from: "n2", to: "n3" }],
    });
    const result = parseContext(buildSelectedContext(selectedElements, graphContext));
    expect(result.neighborNodes).toHaveLength(1);
    expect(result.neighborNodes![0]!.id).toBe("n3");
    expect(result.neighborNodes![0]!.label).toBe("Reject");
  });

  it("should limit neighbor nodes to 10", () => {
    const nodes = [
      { id: "n1", label: "Selected" },
      ...Array.from({ length: 15 }, (_, i) => ({ id: `nb${i}`, label: `Neighbor ${i}` })),
    ];
    const edges = Array.from({ length: 15 }, (_, i) => ({
      id: `e${i}`,
      from: "n1",
      to: `nb${i}`,
    }));
    const graphContext = JSON.stringify({ nodes, edges });
    const result = parseContext(
      buildSelectedContext([{ id: "n1", type: "activity", label: "Selected" }], graphContext),
    );
    expect(result.neighborNodes).toHaveLength(10);
  });

  it("should handle malformed graphContext gracefully", () => {
    const result = parseContext(buildSelectedContext(selectedElements, "not valid json"));
    expect(result.selectedNodes).toEqual(selectedElements);
  });

  it("should handle graphContext with missing nodes/edges arrays", () => {
    const result = parseContext(buildSelectedContext(selectedElements, "{}"));
    expect(result.selectedNodes).toEqual([]);
    expect(result.connectedEdges).toEqual([]);
    expect(result.neighborNodes).toEqual([]);
  });
});
