import { generateObject } from "ai";
import { z } from "zod";

import { modelStrategies } from "../chat/strategies";
import { Model } from "../chat/types";

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

import { TYPE_INFERENCE_RULES } from "./system-prompt";

const typeInferenceSchema = z.object({
  type: z.enum([
    "bpmn",
    "er",
    "orgchart",
    "architecture",
    "sequence",
    "flowchart",
  ]),
  confidence: z.number().min(0).max(1),
});

export type TypeInferenceResult = z.infer<typeof typeInferenceSchema>;

const TYPE_INFERENCE_PROMPT = `Classify this description into a diagram type.

${TYPE_INFERENCE_RULES}

Return the most likely diagram type and your confidence (0-1).`;

export async function inferDiagramType(
  description: string,
): Promise<TypeInferenceResult> {
  const result = await generateObject({
    model: modelStrategies.languageModel(Model.CLAUDE_HAIKU_4_5),
    schema: typeInferenceSchema,
    prompt: `${TYPE_INFERENCE_PROMPT}\n\nDescription: "${description}"`,
  });
  return result.object;
}
