import { describe, expect, it } from "vitest";
import { parseAgency } from "../parser.js";
import { TypescriptPreprocessor } from "../preprocessors/typescriptPreprocessor.js";
import { buildCompilationUnit } from "../compilationUnit.js";
import { AgencyGenerator } from "../backends/agencyGenerator.js";
describe("tag parsing", () => {
    it("parses @goal as a standalone node before a graph node", () => {
        const code = `
@goal("Classify messages")
node main(msg: string): string {
  const result: string = llm("classify: \${msg}")
  return result
}`;
        const result = parseAgency(code);
        expect(result.success).toBe(true);
        if (!result.success)
            return;
        const tagNode = result.result.nodes.find((n) => n.type === "tag");
        expect(tagNode).toBeDefined();
        const tag = tagNode;
        expect(tag.name).toBe("goal");
        expect(tag.arguments).toHaveLength(1);
        expect(tag.arguments[0]).toMatchObject({
            type: "string",
            segments: [{ type: "text", value: "Classify messages" }],
        });
    });
    it("parses @optimize inside a node body", () => {
        const code = `
node main(msg: string): string {
  @optimize
  const result: string = llm("classify: \${msg}")
  return result
}`;
        const result = parseAgency(code);
        expect(result.success).toBe(true);
        if (!result.success)
            return;
        const graphNode = result.result.nodes.find((n) => n.type === "graphNode");
        const tagNode = graphNode.body.find((n) => n.type === "tag");
        expect(tagNode).toBeDefined();
        expect(tagNode.name).toBe("optimize");
    });
});
describe("tag preprocessing", () => {
    it("attaches @goal to the following graph node", () => {
        const code = `
@goal("Classify messages")
node main(msg: string): string {
  const result: string = llm("classify: \${msg}")
  return result
}`;
        const parsed = parseAgency(code);
        expect(parsed.success).toBe(true);
        if (!parsed.success)
            return;
        const info = buildCompilationUnit(parsed.result);
        const preprocessor = new TypescriptPreprocessor(parsed.result, {}, info);
        const processed = preprocessor.preprocess();
        const tagNodes = processed.nodes.filter((n) => n.type === "tag");
        expect(tagNodes).toHaveLength(0);
        const graphNode = processed.nodes.find((n) => n.type === "graphNode");
        expect(graphNode.tags).toHaveLength(1);
        expect(graphNode.tags[0].name).toBe("goal");
        expect(graphNode.tags[0].arguments).toHaveLength(1);
        expect(graphNode.tags[0].arguments[0]).toMatchObject({
            type: "string",
            segments: [{ type: "text", value: "Classify messages" }],
        });
    });
    it("attaches @optimize to the following assignment inside a body", () => {
        const code = `
node main(msg: string): string {
  @optimize
  const result: string = llm("classify: \${msg}")
  return result
}`;
        const parsed = parseAgency(code);
        expect(parsed.success).toBe(true);
        if (!parsed.success)
            return;
        const info = buildCompilationUnit(parsed.result);
        const preprocessor = new TypescriptPreprocessor(parsed.result, {}, info);
        const processed = preprocessor.preprocess();
        const graphNode = processed.nodes.find((n) => n.type === "graphNode");
        const tagNodes = graphNode.body.filter((n) => n.type === "tag");
        expect(tagNodes).toHaveLength(0);
        const assignment = graphNode.body.find((n) => n.type === "assignment");
        expect(assignment.tags).toHaveLength(1);
        expect(assignment.tags[0].name).toBe("optimize");
    });
});
describe("tag formatting", () => {
    it("round-trips tags through parse and format (standalone path)", () => {
        const code = `@goal("Classify messages accurately")
node main(msg: string): string {
  @optimize
  const result: string = llm("classify: \${msg}")
  return result
}`;
        const result = parseAgency(code);
        expect(result.success).toBe(true);
        if (!result.success)
            return;
        const generator = new AgencyGenerator();
        const output = generator.generate(result.result);
        expect(output.output).toContain(`@goal("Classify messages accurately")\nnode main`);
        expect(output.output).toContain(`  @optimize\n  const result: string`);
    });
    it("formats tags attached to nodes (post-preprocessing)", () => {
        const code = `@goal("Classify messages accurately")
node main(msg: string): string {
  @optimize
  const result: string = llm("classify: \${msg}")
  return result
}`;
        const parsed = parseAgency(code);
        expect(parsed.success).toBe(true);
        if (!parsed.success)
            return;
        const info = buildCompilationUnit(parsed.result);
        const preprocessor = new TypescriptPreprocessor(parsed.result, {}, info);
        const processed = preprocessor.preprocess();
        const generator = new AgencyGenerator();
        const output = generator.generate(processed);
        expect(output.output).toContain(`@goal("Classify messages accurately")\nnode main`);
        expect(output.output).toContain(`  @optimize\n  const result: string`);
    });
    it("formats tag with no arguments", () => {
        const code = `@optimize
node main(): string {
  const result: string = llm("hello")
  return result
}`;
        const result = parseAgency(code);
        expect(result.success).toBe(true);
        if (!result.success)
            return;
        const generator = new AgencyGenerator();
        const output = generator.generate(result.result);
        expect(output.output).toContain(`@optimize\nnode main`);
    });
    it("formats tag with multiple arguments including spaces", () => {
        const code = `@goal("classify messages", "with high accuracy")
node main(): string {
  const result: string = llm("hello")
  return result
}`;
        const result = parseAgency(code);
        expect(result.success).toBe(true);
        if (!result.success)
            return;
        const generator = new AgencyGenerator();
        const output = generator.generate(result.result);
        expect(output.output).toContain(`@goal("classify messages", "with high accuracy")\nnode main`);
    });
    it("formats bare identifier tag arguments without quotes", () => {
        const code = `@goal(fast)
node main(): string {
  const result: string = llm("hello")
  return result
}`;
        const result = parseAgency(code);
        expect(result.success).toBe(true);
        if (!result.success)
            return;
        const generator = new AgencyGenerator();
        const output = generator.generate(result.result);
        expect(output.output).toContain(`@goal(fast)\nnode main`);
    });
});
