import { Ollama } from 'ollama';
import z from 'zod';
import { repairJson } from '@toolsycc/json-repair';
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: 'user' | 'assistant' | 'system'; content: string } {
    return {
      role: msg.role === 'system' ? 'system' : msg.role === 'user' ? 'user' : 'assistant',
      content: msg.content,
    };
  }

  async generateObject<T>(input: {
    schema: z.ZodTypeAny;
    messages: Message[];
  }): Promise<T> {
    const messages = input.messages.map((m) => this.toOllama(m));
    const res = await this.client.chat({
      model: this.model,
      messages,
      format: z.toJSONSchema(input.schema),
    });
    const content = res.message?.content;
    if (!content) throw new Error('No response from Ollama');
    const parsed = JSON.parse(
      repairJson(content, { extractJson: true }) as string,
    );
    return input.schema.parse(parsed) as T;
  }
}
