inkling/src/main/ai/LocalOllamaProvider.ts

import { request } from 'undici';
import { parseAiResponse, type AiResponse } from './schema.js';
import { buildPrompt } from './prompt.js';
import { buildVisionPrompt } from './visionPrompt.js';
import type { GenerateInput, GenerateOptions, HealthResult, InferenceProvider } from './InferenceProvider.js';
import { DEFAULT_OLLAMA_ENDPOINT, DEFAULT_OLLAMA_MODEL } from '../../shared/constants.js';

function classifyFetchError(e: unknown): 'network' | 'timeout' | 'dns' | 'other' {
  const msg = ((e as Error)?.message ?? '').toLowerCase();
  if (msg.includes('aborted') || msg.includes('timeout')) return 'timeout';
  if (msg.includes('econnrefused') || msg.includes('econnreset')) return 'network';
  if (msg.includes('enotfound') || msg.includes('eai_again')) return 'dns';
  return 'other';
}

export interface LocalOllamaOptions {
  endpoint?: string;
  model?: string;
  timeoutMs?: number;
  temperature?: number;
  numPredict?: number;
}

export class LocalOllamaProvider implements InferenceProvider {
  readonly name: string;
  private endpoint: string;
  private model: string;
  private timeoutMs: number;
  private temperature: number;
  private numPredict: number;
  private abortController: AbortController | null = null;

  constructor(opts: LocalOllamaOptions = {}) {
    this.endpoint = opts.endpoint ?? DEFAULT_OLLAMA_ENDPOINT;
    this.model = opts.model ?? DEFAULT_OLLAMA_MODEL;
    this.timeoutMs = opts.timeoutMs ?? 120_000;
    this.temperature = opts.temperature ?? 0.2;
    this.numPredict = opts.numPredict ?? 512;
    this.name = `local-ollama/${this.model}`;
  }

  async generate(input: GenerateInput, opts?: GenerateOptions): Promise<AiResponse> {
    const useVision = !!opts?.visionModel && (input.images?.length ?? 0) > 0;
    const model = useVision ? opts!.visionModel! : this.model;
    const prompt = useVision
      ? buildVisionPrompt(input.text, input.todayKst, input.dueDateCandidates.map((c) => c.iso ?? c.matchedToken ?? ''), input.vocab ?? [])
      : buildPrompt(input.text, input.todayKst, input.dueDateCandidates, input.vocab ?? []);

    this.abortController = new AbortController();
    const timer = setTimeout(() => this.abortController?.abort(), this.timeoutMs);
    try {
      const body: Record<string, unknown> = {
        model,
        prompt,
        format: 'json',
        stream: false,
        options: { temperature: this.temperature, num_predict: this.numPredict }
      };
      if (useVision) {
        body.images = input.images!.map((i) => i.base64);
      }
      const res = await request(`${this.endpoint}/api/generate`, {
        method: 'POST',
        headers: { 'content-type': 'application/json' },
        body: JSON.stringify(body),
        signal: this.abortController.signal
      });
      if (res.statusCode < 200 || res.statusCode >= 300) {
        throw new Error(`ollama http ${res.statusCode}`);
      }
      const responseBody = (await res.body.json()) as { response?: string };
      if (!responseBody.response) throw new Error('missing response field');
      let parsed: unknown;
      try { parsed = JSON.parse(responseBody.response); }
      catch (err) { throw new Error(`invalid json in response: ${String(err)}`); }
      return parseAiResponse(parsed);
    } finally {
      clearTimeout(timer);
      this.abortController = null;
    }
  }

  /** v0.2.3.1 — 외부에서 in-flight generate 강제 중단. ProviderHolder.replace 시 사용. */
  abort(): void {
    this.abortController?.abort();
  }

  /**
   * v0.2.9 Cut B Task 9 — raw JSON 호출 (classifyStatus 등 자체 prompt 용).
   * `format: 'json'` + `stream: false` 로 Ollama 가 valid JSON 문자열을 반환하도록 강제.
   * abortController / timeout 은 generate() 와 동일 패턴.
   */
  async generateRaw(prompt: string): Promise<string> {
    this.abortController = new AbortController();
    const timer = setTimeout(() => this.abortController?.abort(), this.timeoutMs);
    try {
      const res = await request(`${this.endpoint}/api/generate`, {
        method: 'POST',
        headers: { 'content-type': 'application/json' },
        body: JSON.stringify({
          model: this.model,
          prompt,
          format: 'json',
          stream: false,
          options: { temperature: this.temperature, num_predict: this.numPredict }
        }),
        signal: this.abortController.signal
      });
      if (res.statusCode < 200 || res.statusCode >= 300) {
        throw new Error(`ollama http ${res.statusCode}`);
      }
      const body = (await res.body.json()) as { response?: string };
      if (!body.response) throw new Error('missing response field');
      return body.response;
    } finally {
      clearTimeout(timer);
      this.abortController = null;
    }
  }

  async healthCheck(): Promise<HealthResult> {
    try {
      const res = await request(`${this.endpoint}/api/tags`, { method: 'GET' });
      if (res.statusCode !== 200) return { ok: false, reason: `tags http ${res.statusCode}` };
      const body = (await res.body.json()) as { models?: Array<{ name: string }> };
      const found = body.models?.some((m) => m.name === this.model);
      return found ? { ok: true, model: this.model }
                   : { ok: false, reason: `${this.model} not installed` };
    } catch (err) {
      const cls = classifyFetchError(err);
      return { ok: false, reason: `unreachable:${cls}` };
    }
  }
}