"""
Ollama Client — Local LLM inference
=====================================
Real implementation against the Ollama REST API (http://localhost:11434).
No cloud APIs, no API keys. Runs entirely air-gapped.

Ollama API:
  POST /api/chat          — multi-turn chat (OpenAI-compatible messages)
  POST /api/generate      — single completion
  GET  /api/tags          — list installed models
  POST /api/pull          — pull a model

This client exposes the same interface as GroqClient so Pipeline can
swap them transparently via settings.backend = "ollama".
"""
from __future__ import annotations

import json
import asyncio
from typing import AsyncIterator, Optional

import httpx

from forge.core.api import strip_fences, extract_json, APIError


class OllamaClient:
    """
    Drop-in replacement for GroqClient + NvidiaClient when running locally.
    Talks directly to a running `ollama serve` process.
    """

    DEFAULT_HOST = "http://localhost:11434"

    def __init__(
        self,
        host: str = DEFAULT_HOST,
        timeout: int = 300,          # local can be slow on CPU
        default_model: str = "codellama:70b",
    ):
        self.host = host.rstrip("/")
        self.default_model = default_model
        self._timeout = httpx.Timeout(timeout, connect=10.0)

    # ── Health / discovery ────────────────────────────────────────────────────

    async def is_available(self) -> bool:
        """Return True if Ollama is running and reachable."""
        try:
            async with httpx.AsyncClient(timeout=httpx.Timeout(5.0)) as client:
                resp = await client.get(f"{self.host}/api/tags")
                return resp.status_code == 200
        except Exception:
            return False

    async def list_models(self) -> list[str]:
        """Return names of all locally installed models."""
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            resp = await client.get(f"{self.host}/api/tags")
            if not resp.is_success:
                return []
            data = resp.json()
            return [m["name"] for m in data.get("models", [])]

    async def pull_model(self, model: str) -> bool:
        """
        Pull a model from the Ollama registry if not present.
        Streams progress lines — returns True on success.
        """
        async with httpx.AsyncClient(timeout=httpx.Timeout(3600, connect=10.0)) as client:
            async with client.stream(
                "POST",
                f"{self.host}/api/pull",
                json={"name": model, "stream": True},
            ) as resp:
                if not resp.is_success:
                    return False
                async for line in resp.aiter_lines():
                    if not line:
                        continue
                    try:
                        data = json.loads(line)
                        if data.get("status") == "success":
                            return True
                    except Exception:
                        pass
        return False

    async def ensure_model(self, model: str) -> bool:
        """Pull model if not installed. Returns True if model is available after."""
        installed = await self.list_models()
        # Normalize: 'codellama:70b' matches 'codellama:70b-instruct' prefix check
        if any(m.startswith(model.split(":")[0]) for m in installed):
            return True
        return await self.pull_model(model)

    # ── Completion — same signature as GroqClient.complete ────────────────────

    async def complete(
        self,
        messages: list[dict],
        model: Optional[str] = None,
        temperature: float = 0.3,
        max_tokens: int = 8000,
    ) -> str:
        """Multi-turn chat completion. Returns full response text."""
        model = model or self.default_model
        payload = {
            "model": model,
            "messages": messages,
            "stream": False,
            "options": {
                "temperature": temperature,
                "num_predict": max_tokens,
                "num_ctx": 16384,        # context window
            },
        }
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            resp = await client.post(f"{self.host}/api/chat", json=payload)
            if not resp.is_success:
                raise APIError("Ollama", resp.status_code, resp.text[:200])
            data = resp.json()
            return data["message"]["content"]

    async def complete_json(
        self,
        messages: list[dict],
        model: Optional[str] = None,
        temperature: float = 0.2,
    ) -> dict:
        raw = await self.complete(messages, model=model, temperature=temperature)
        return extract_json(raw)

    # ── Code generation — same signature as NvidiaClient.generate_code ───────

    async def generate_code(
        self,
        system_prompt: str,
        user_prompt: str,
        model: Optional[str] = None,
        max_tokens: int = 4096,
    ) -> str:
        """Single-turn code generation."""
        model = model or self.default_model
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user",   "content": user_prompt},
        ]
        raw = await self.complete(messages, model=model, temperature=0.1, max_tokens=max_tokens)
        return strip_fences(raw)

    async def generate_code_with_fallback(
        self,
        system_prompt: str,
        user_prompt: str,
        model: str,
        groq_client=None,   # ignored — Ollama IS the fallback
        groq_model: str = "",
    ) -> tuple[str, str]:
        """Compatible interface with NvidiaClient. Always uses local Ollama."""
        code = await self.generate_code(system_prompt, user_prompt, model)
        return code, "ollama"

    # ── Streaming ─────────────────────────────────────────────────────────────

    async def stream(
        self,
        messages: list[dict],
        model: Optional[str] = None,
        temperature: float = 0.2,
        max_tokens: int = 6000,
    ) -> AsyncIterator[str]:
        model = model or self.default_model
        payload = {
            "model": model,
            "messages": messages,
            "stream": True,
            "options": {"temperature": temperature, "num_predict": max_tokens},
        }
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            async with client.stream(
                "POST", f"{self.host}/api/chat", json=payload
            ) as resp:
                if not resp.is_success:
                    raise APIError("Ollama", resp.status_code, await resp.aread())
                async for line in resp.aiter_lines():
                    if not line:
                        continue
                    try:
                        chunk = json.loads(line)
                        token = chunk.get("message", {}).get("content", "")
                        if token:
                            yield token
                        if chunk.get("done"):
                            break
                    except Exception:
                        pass
