export interface OllamaModelInfo {
  name: string;
  modified_at: string;
  size: number;
  digest: string;
  details?: {
    format?: string;
    family?: string;
    families?: string[] | null;
    parameter_size?: string;
    quantization_level?: string;
  };
}

export interface OllamaToolCall {
  id?: string;
  function: {
    name: string;
    arguments: Record<string, unknown> | string;
  };
}

export interface OllamaChatMessage {
  role: "system" | "user" | "assistant" | "tool";
  content: string;
  tool_calls?: OllamaToolCall[];
  tool_call_id?: string;
  name?: string;
}

export interface OllamaToolDefinition {
  type: "function";
  function: {
    name: string;
    description: string;
    parameters: {
      type: "object";
      properties: Record<string, unknown>;
      required?: string[];
    };
  };
}

export interface OllamaChatChunk {
  model: string;
  created_at: string;
  message?: {
    role: string;
    content: string;
    thinking?: string;
    tool_calls?: OllamaToolCall[];
  };
  done: boolean;
  done_reason?: string;
  total_duration?: number; // ns
  load_duration?: number; // ns
  prompt_eval_count?: number;
  prompt_eval_duration?: number; // ns
  eval_count?: number;
  eval_duration?: number; // ns
}

export interface OllamaPullChunk {
  status: string;
  digest?: string;
  total?: number;
  completed?: number;
  error?: string;
}

export interface OllamaShowResponse {
  modelfile?: string;
  parameters?: string;
  template?: string;
  details?: {
    parameter_size?: string;
    quantization_level?: string;
    family?: string;
    families?: string[] | null;
  };
  capabilities?: string[];
}

// In-memory cache; capabilities are static per model name.
const showCache = new Map<string, OllamaShowResponse>();

export async function showModel(baseUrl: string, name: string): Promise<OllamaShowResponse> {
  const key = `${ensureUrl(baseUrl)}::${name}`;
  const cached = showCache.get(key);
  if (cached) return cached;
  const res = await fetch(`${ensureUrl(baseUrl)}/api/show`, {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({ name }),
  });
  if (!res.ok) throw new Error(`Ollama show failed: ${res.status}`);
  const data = (await res.json()) as OllamaShowResponse;
  showCache.set(key, data);
  return data;
}

export function supportsThinking(show: OllamaShowResponse | null | undefined): boolean {
  if (!show) return false;
  const caps = show.capabilities ?? [];
  return caps.includes("thinking");
}

/**
 * Whitelist of model name prefixes known to support tool calling.
 * Used as fallback when Ollama doesn't report capabilities.
 */
const TOOLS_WHITELIST_PREFIXES = [
  "qwen2.5",
  "qwen3",
  "llama3.1",
  "llama3.2",
  "llama3.3",
  "mistral-nemo",
  "mistral-small",
  "mistral-large",
  "command-r",
  "command-r-plus",
  "hermes3",
  "firefunction-v2",
];

export function supportsTools(
  modelName: string,
  show: OllamaShowResponse | null | undefined,
): boolean {
  if (show?.capabilities) return show.capabilities.includes("tools");
  // Fallback: name-prefix whitelist
  const normalized = modelName.toLowerCase().split(":")[0] ?? "";
  return TOOLS_WHITELIST_PREFIXES.some((p) => normalized.startsWith(p));
}

function ensureUrl(url: string): string {
  return url.replace(/\/+$/, "");
}

export async function listModels(baseUrl: string): Promise<OllamaModelInfo[]> {
  const res = await fetch(`${ensureUrl(baseUrl)}/api/tags`, { cache: "no-store" });
  if (!res.ok) throw new Error(`Ollama list failed: ${res.status}`);
  const data = (await res.json()) as { models?: OllamaModelInfo[] };
  return data.models ?? [];
}

export async function testConnection(baseUrl: string): Promise<{ ok: boolean; version?: string; error?: string }> {
  try {
    const res = await fetch(`${ensureUrl(baseUrl)}/api/version`, {
      cache: "no-store",
      signal: AbortSignal.timeout(5000),
    });
    if (!res.ok) return { ok: false, error: `HTTP ${res.status}` };
    const data = (await res.json()) as { version?: string };
    return { ok: true, version: data.version };
  } catch (e) {
    return { ok: false, error: e instanceof Error ? e.message : "Unknown error" };
  }
}

export async function deleteModel(baseUrl: string, name: string): Promise<void> {
  const res = await fetch(`${ensureUrl(baseUrl)}/api/delete`, {
    method: "DELETE",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({ name }),
  });
  if (!res.ok) throw new Error(`Ollama delete failed: ${res.status}`);
}

export async function pullModelStream(baseUrl: string, name: string, signal?: AbortSignal): Promise<ReadableStream<Uint8Array>> {
  const res = await fetch(`${ensureUrl(baseUrl)}/api/pull`, {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({ name, stream: true }),
    signal,
  });
  if (!res.ok || !res.body) throw new Error(`Ollama pull failed: ${res.status}`);
  return res.body;
}

export async function chatStream(
  baseUrl: string,
  body: {
    model: string;
    messages: OllamaChatMessage[];
    options?: Record<string, unknown>;
    think?: boolean;
    keep_alive?: number | string;
    tools?: OllamaToolDefinition[];
  },
  signal?: AbortSignal,
): Promise<ReadableStream<Uint8Array>> {
  const res = await fetch(`${ensureUrl(baseUrl)}/api/chat`, {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({ ...body, stream: true }),
    signal,
  });
  if (!res.ok || !res.body) {
    const text = await res.text().catch(() => "");
    throw new Error(`Ollama chat failed: ${res.status} ${text}`);
  }
  return res.body;
}

export async function chatOnce(
  baseUrl: string,
  body: { model: string; messages: OllamaChatMessage[]; options?: Record<string, unknown> },
): Promise<string> {
  const res = await fetch(`${ensureUrl(baseUrl)}/api/chat`, {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({ ...body, stream: false }),
  });
  if (!res.ok) throw new Error(`Ollama chat failed: ${res.status}`);
  const data = (await res.json()) as { message?: { content?: string } };
  return data.message?.content ?? "";
}

/**
 * Parse Ollama NDJSON stream into typed chunks.
 */
export async function* parseNdjson<T>(stream: ReadableStream<Uint8Array>): AsyncGenerator<T> {
  const reader = stream.getReader();
  const decoder = new TextDecoder();
  let buffer = "";
  try {
    while (true) {
      const { value, done } = await reader.read();
      if (done) break;
      buffer += decoder.decode(value, { stream: true });
      let nl: number;
      while ((nl = buffer.indexOf("\n")) >= 0) {
        const line = buffer.slice(0, nl).trim();
        buffer = buffer.slice(nl + 1);
        if (!line) continue;
        try {
          yield JSON.parse(line) as T;
        } catch {
          // skip malformed line
        }
      }
    }
    const tail = buffer.trim();
    if (tail) {
      try {
        yield JSON.parse(tail) as T;
      } catch {
        /* ignore */
      }
    }
  } finally {
    reader.releaseLock();
  }
}

/**
 * Split assistant content into visible content + thinking by parsing <think>...</think> tags.
 */
export function splitThinking(content: string): { content: string; thinking: string } {
  let thinking = "";
  const cleaned = content.replace(/<think>([\s\S]*?)<\/think>/g, (_, inner: string) => {
    thinking += (thinking ? "\n\n" : "") + inner.trim();
    return "";
  });
  return { content: cleaned.trim(), thinking };
}