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

// Mock the AI SDK generateObject before importing the module
vi.mock("ai", () => ({
  generateObject: vi.fn(),
}));

// Mock model strategies
vi.mock("../chat/strategies", () => ({
  modelStrategies: {
    languageModel: vi.fn(() => "mock-model"),
  },
}));

import { generateObject } from "ai";
import { inferDiagramType } from "./type-inference";

const mockGenerateObject = vi.mocked(generateObject);

describe("inferDiagramType", () => {
  it("should return the inferred type and confidence for ER description", async () => {
    mockGenerateObject.mockResolvedValueOnce({
      object: { type: "er", confidence: 0.95 },
    } as never);

    const result = await inferDiagramType("database schema for user management");
    expect(result.type).toBe("er");
    expect(result.confidence).toBe(0.95);
  });

  it("should return the inferred type for BPMN description", async () => {
    mockGenerateObject.mockResolvedValueOnce({
      object: { type: "bpmn", confidence: 0.88 },
    } as never);

    const result = await inferDiagramType("order approval workflow with manager review");
    expect(result.type).toBe("bpmn");
    expect(result.confidence).toBeGreaterThan(0);
  });

  it("should return the inferred type for architecture description", async () => {
    mockGenerateObject.mockResolvedValueOnce({
      object: { type: "architecture", confidence: 0.92 },
    } as never);

    const result = await inferDiagramType("microservices system design with API gateway");
    expect(result.type).toBe("architecture");
  });

  it("should return the inferred type for orgchart description", async () => {
    mockGenerateObject.mockResolvedValueOnce({
      object: { type: "orgchart", confidence: 0.9 },
    } as never);

    const result = await inferDiagramType("company team structure with reporting lines");
    expect(result.type).toBe("orgchart");
  });

  it("should return the inferred type for sequence description", async () => {
    mockGenerateObject.mockResolvedValueOnce({
      object: { type: "sequence", confidence: 0.85 },
    } as never);

    const result = await inferDiagramType("API call flow between browser and server");
    expect(result.type).toBe("sequence");
  });

  it("should return the inferred type for flowchart description", async () => {
    mockGenerateObject.mockResolvedValueOnce({
      object: { type: "flowchart", confidence: 0.87 },
    } as never);

    const result = await inferDiagramType("decision tree for loan approval");
    expect(result.type).toBe("flowchart");
  });

  it("should pass the description in the prompt", async () => {
    mockGenerateObject.mockResolvedValueOnce({
      object: { type: "er", confidence: 0.9 },
    } as never);

    const description = "user registration database tables";
    await inferDiagramType(description);

    expect(mockGenerateObject).toHaveBeenCalledWith(
      expect.objectContaining({
        prompt: expect.stringContaining(description) as string,
      }),
    );
  });

  it("should include type inference rules in the prompt", async () => {
    mockGenerateObject.mockResolvedValueOnce({
      object: { type: "bpmn", confidence: 0.9 },
    } as never);

    await inferDiagramType("some description");

    expect(mockGenerateObject).toHaveBeenCalledWith(
      expect.objectContaining({
        prompt: expect.stringContaining("Business processes") as string,
      }),
    );
  });

  it("should use a schema that covers all 6 diagram types", async () => {
    mockGenerateObject.mockResolvedValueOnce({
      object: { type: "er", confidence: 0.9 },
    } as never);

    await inferDiagramType("test");

    const call = mockGenerateObject.mock.calls[0]![0] as { schema: { shape: { type: { options: string[] } } } };
    const typeOptions = call.schema.shape.type.options as string[];
    expect(typeOptions).toEqual(
      expect.arrayContaining(["bpmn", "er", "orgchart", "architecture", "sequence", "flowchart"]),
    );
    expect(typeOptions).toHaveLength(6);
  });

  it("should propagate errors from generateObject", async () => {
    mockGenerateObject.mockRejectedValueOnce(new Error("API error"));

    await expect(inferDiagramType("test description")).rejects.toThrow("API error");
  });
});
