import { NextRequest } from 'next/server';
import { OLLAMA_URL } from '../../../lib/ollama-config';
import {
  LlmRouter,
  createLocalOllamaProvider,
  type ChatCompletionRequest,
  type RoutePlan,
} from '../../../lib/llm-router';

interface OllamaTagsResponse {
  models?: Array<{ name?: string }>;
}

interface OllamaPsResponse {
  models?: Array<{ name?: string }>;
}

async function fetchOllamaJson<T>(path: string): Promise<T> {
  const response = await fetch(`${OLLAMA_URL}${path}`, {
    cache: 'no-store',
    signal: AbortSignal.timeout(5000),
  });

  if (!response.ok) {
    throw new Error(`Ollama ${path} failed: ${response.status}`);
  }

  return (await response.json()) as T;
}

async function resolvePlan(request: ChatCompletionRequest): Promise<RoutePlan> {
  const [tags, ps] = await Promise.all([
    fetchOllamaJson<OllamaTagsResponse>('/api/tags'),
    fetchOllamaJson<OllamaPsResponse>('/api/ps').catch(() => ({ models: [] })),
  ]);

  const installedModels = (tags.models ?? [])
    .map(model => model.name?.trim() || '')
    .filter(Boolean);
  const runningModels = new Set(
    (ps.models ?? []).map(model => model.name?.trim() || '').filter(Boolean)
  );

  if (installedModels.length === 0) {
    throw new Error('No local Ollama models installed');
  }

  // Put currently loaded models first so the router prefers them when scores are equal.
  const orderedModels = [
    ...installedModels.filter(model => runningModels.has(model)),
    ...installedModels.filter(model => !runningModels.has(model)),
  ];

  const router = new LlmRouter({
    providers: [createLocalOllamaProvider(orderedModels, `${OLLAMA_URL}/v1`)],
    maxRetries: 1,
  });

  return router.plan(request);
}

export async function POST(request: NextRequest) {
  try {
    const body = await request.json();
    const { model, messages } = body;

    if (!model || !Array.isArray(messages)) {
      return new Response(JSON.stringify({ error: 'Missing model or messages' }), {
        status: 400,
        headers: { 'Content-Type': 'application/json' },
      });
    }

    const plan = await resolvePlan({
      messages,
      model: model && model !== '__auto__' ? `local-ollama:${String(model)}` : undefined,
    });

    const response = await fetch(`${OLLAMA_URL}/api/chat`, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({ model: plan.model.id, messages, stream: true }),
    });

    if (!response.ok || !response.body) {
      return new Response(JSON.stringify({ error: `Ollama chat error: ${response.status}` }), {
        status: 500,
        headers: { 'Content-Type': 'application/json' },
      });
    }

    return new Response(response.body, {
      headers: {
        'Content-Type': 'application/x-ndjson',
        'Transfer-Encoding': 'chunked',
        'Cache-Control': 'no-cache',
        'x-by-provider': plan.provider.name,
        'x-by-model': plan.model.id,
        'x-by-category': plan.category,
      },
    });
  } catch (err) {
    return new Response(JSON.stringify({ error: String(err) }), {
      status: 500,
      headers: { 'Content-Type': 'application/json' },
    });
  }
}
