import OpenAI from 'openai';
import z from 'zod';
import { BaseLLM } from './base/llm.js';
import type { Message, ToolMessage } 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 === 'tool') {
      return {
        role: 'tool',
        tool_call_id: msg.id,
        content: msg.content,
      };
    }
    if (msg.role === 'assistant') {
      const m = msg as { content: string; tool_calls?: { id: string; name: string; arguments: Record<string, unknown> }[] };
      return {
        role: 'assistant',
        content: m.content,
        tool_calls: m.tool_calls?.map((tc) => ({
          id: tc.id,
          type: 'function' as const,
          function: { name: tc.name, arguments: JSON.stringify(tc.arguments) },
        })),
      };
    }
    return { role: 'user', content: msg.content };
  }

  async generateText(input: { messages: Message[]; tools?: { name: string; description: string; schema: z.ZodObject<any> }[] }): Promise<{ content: string; toolCalls: { id: string; name: string; arguments: Record<string, unknown> }[] }> {
    const tools = input.tools?.map((t) => ({
      type: 'function' as const,
      function: {
        name: t.name,
        description: t.description,
        parameters: z.toJSONSchema(t.schema),
      },
    }));

    const res = await this.client.chat.completions.create({
      model: this.model,
      messages: input.messages.map((m) => this.toOpenAI(m)),
      tools: tools?.length ? tools : undefined,
      temperature: 0.7,
      max_tokens: 4096,
    });

    const msg = res.choices?.[0]?.message;
    if (!msg) throw new Error('No response from OpenAI');

    const toolCalls = (msg.tool_calls ?? []).map((tc) => {
      const fn = 'function' in tc ? (tc as { function?: { name?: string; arguments?: string } }).function : null;
      return {
        id: tc.id ?? crypto.randomUUID(),
        name: fn?.name ?? 'unknown',
        arguments: (() => {
          try {
            return JSON.parse(fn?.arguments ?? '{}') as Record<string, unknown>;
          } catch {
            return {};
          }
        })(),
      };
    });

    return {
      content: msg.content ?? '',
      toolCalls,
    };
  }
}
