import OpenAI from 'openai';
import { zodResponseFormat } from 'openai/helpers/zod';
import { repairJson } from '@toolsycc/json-repair';
import { BaseLLM } from './base/llm.js';
import type { Message } from '../types.js';

type OpenAIConfig = {
  apiKey: string;
  model: string;
  baseURL?: string;
};

export class OpenAILLM extends BaseLLM {
  private client: OpenAI;
  private model: string;

  constructor(config: OpenAIConfig) {
    super();
    this.model = config.model;
    this.client = new OpenAI({
      apiKey: config.apiKey,
      baseURL: config.baseURL ?? 'https://api.openai.com/v1',
    });
  }

  private toOpenAI(msg: Message): OpenAI.Chat.Completions.ChatCompletionMessageParam {
    if (msg.role === 'system') return { role: 'system', content: msg.content };
    if (msg.role === 'assistant') return { role: 'assistant', content: msg.content };
    return { role: 'user', content: msg.content };
  }

  async generateObject<T>(input: {
    schema: Parameters<typeof zodResponseFormat>[0];
    messages: Message[];
  }): Promise<T> {
    const messages = input.messages.map((m) => this.toOpenAI(m));
    const res = await this.client.chat.completions.parse({
      model: this.model,
      messages,
      response_format: zodResponseFormat(input.schema, 'object'),
    });
    const content = res.choices?.[0]?.message?.content;
    if (!content) throw new Error('No response from OpenAI');
    const parsed = JSON.parse(
      repairJson(content, { extractJson: true }) as string,
    );
    return input.schema.parse(parsed) as T;
  }
}
