import { Ollama } from 'ollama';
import z from 'zod';
import crypto from 'node:crypto';
import { BaseLLM } from './base/llm.js';
import type { Message } from '../types.js';

type OllamaConfig = { baseURL: string; model: string };

export class OllamaLLM extends BaseLLM {
  private client: Ollama;
  private model: string;

  constructor(config: OllamaConfig) {
    super();
    this.model = config.model;
    this.client = new Ollama({
      host: config.baseURL || 'http://localhost:11434',
    });
  }

  private toOllama(msg: Message): { role: string; content: string; tool_calls?: { id: string; type: string; function: { name: string; arguments: unknown } }[] } {
    if (msg.role === 'tool') return { role: 'tool', content: msg.content } as { role: string; content: string };
    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',
          function: { name: tc.name, arguments: tc.arguments },
        })),
      };
    }
    return { role: msg.role, content: (msg as { content: string }).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 ollamaTools = input.tools?.map((t) => ({
      type: 'function' as const,
      function: {
        name: t.name,
        description: t.description,
        parameters: (z.toJSONSchema(t.schema) as Record<string, unknown>).properties ?? {},
      },
    }));

    const res = await this.client.chat({
      model: this.model,
      messages: input.messages.map((m) => this.toOllama(m)) as Parameters<typeof this.client.chat>[0]['messages'],
      tools: ollamaTools?.length ? ollamaTools : undefined,
      options: { temperature: 0.7, num_predict: 4096 },
    });

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

    return {
      content: res.message.content ?? '',
      toolCalls,
    };
  }
}
