import z from 'zod';
import type { LlmClient } from '../llm-client.js';
import { getClassifierPrompt } from '../prompts/classifier.js';
import formatChatHistoryAsString from '../utils/formatHistory.js';

const schema = z.object({
  classification: z.object({
    skipSearch: z.boolean(),
    personalSearch: z.boolean(),
    academicSearch: z.boolean(),
    discussionSearch: z.boolean(),
    showWeatherWidget: z.boolean(),
    showStockWidget: z.boolean(),
    showCalculationWidget: z.boolean(),
  }),
  standaloneFollowUp: z.string(),
});

export type ClassifierInput = {
  chatHistory: { role: string; content: string }[];
  query: string;
  llm: LlmClient;
  locale?: string;
  enabledSources: ('web' | 'discussions' | 'academic')[];
};

export async function classify(input: ClassifierInput): Promise<z.infer<typeof schema>> {
  const output = await input.llm.generateObject<z.infer<typeof schema>>({
    messages: [
      { role: 'system', content: getClassifierPrompt(input.locale) },
      {
        role: 'user',
        content: `<conversation_history>\n${formatChatHistoryAsString(input.chatHistory)}\n</conversation_history>\n<user_query>\n${input.query}\n</user_query>`,
      },
    ],
    schema,
  });
  return output;
}