"""
cortex/llm/ollama_provider.py — Ollama Local Inference Implementation of LLMProvider.

Wraps Ollama's HTTP API for fully local, zero-cost, privacy-first LLM inference.
This is the fallback provider when API keys are unavailable or when data must
not leave the machine.

Configuration:
    OLLAMA_BASE_URL — Base URL for Ollama server (default: http://localhost:11434).
    No API key required — Ollama runs locally.

Supported models:
    Completion: Any model available in the local Ollama instance
                (llama3, mistral, phi, gemma, etc.)
    Embedding: nomic-embed-text (768d, recommended), mxbai-embed-large

Limitations vs. cloud providers:
    - Limited context length (depends on model, typically 4K-128K)
    - No guaranteed JSON schema enforcement (relies on model compliance)
    - Slower inference on CPU-only machines
    - Embedding quality may be lower than Voyage/OpenAI

Key implementation details:
    - Uses Ollama's /api/chat endpoint for completions
    - Uses /api/embeddings endpoint for vector generation
    - Cost estimation always returns 0.0 (local inference is free)
    - No Voyage delegation — everything runs through Ollama locally

References:
    https://ollama.com/
    https://github.com/ollama/ollama/blob/main/docs/api.md

Dependencies: httpx (async HTTP client).
Implements: cortex.llm.provider.LLMProvider ABC.
"""

from __future__ import annotations

import json
import logging
import os
import time
from typing import Any

from cortex.llm.provider import (
    CompletionResponse,
    EmbedResponse,
    ExtractionResponse,
    LLMProvider,
    UsageRecord,
    retry_with_backoff,
)

logger = logging.getLogger(__name__)

_DEFAULT_BASE_URL = "http://localhost:11434"


class OllamaProvider(LLMProvider):
    """
    LLMProvider backed by Ollama's local HTTP API.

    No API key needed. Zero external cost. Ideal for local dev + fallback.

    Environment variables:
        OLLAMA_BASE_URL  — optional override, default http://localhost:11434
        OLLAMA_COMPLETION_MODEL — optional model override
        OLLAMA_EMBED_MODEL — optional embed model override

    Supported models (must be pulled via `ollama pull <model>`):
        Completion: llama3.2, mistral, gemma3, deepseek-r1:8b, etc.
        Embedding: nomic-embed-text, mxbai-embed-large, all-minilm
    """

    DEFAULT_COMPLETION_MODEL = "llama3.2"
    DEFAULT_EMBED_MODEL = "nomic-embed-text"
    PROVIDER_NAME = "ollama"

    # Prefixes that indicate a non-Ollama embedding model (OpenAI / Voyage)
    _INCOMPATIBLE_EMBED_PREFIXES = ("text-embedding-", "voyage-")

    def __init__(
        self,
        completion_model: str | None = None,
        embed_model: str | None = None,
        base_url: str | None = None,
    ) -> None:
        super().__init__()
        self._completion_model = (
            completion_model
            or os.environ.get("OLLAMA_COMPLETION_MODEL")
            or self.DEFAULT_COMPLETION_MODEL
        )
        # Validate embed_model — reject OpenAI/Voyage models that Ollama can't use
        candidate_embed = (
            embed_model
            or os.environ.get("OLLAMA_EMBED_MODEL")
            or self.DEFAULT_EMBED_MODEL
        )
        if any(candidate_embed.startswith(p) for p in self._INCOMPATIBLE_EMBED_PREFIXES):
            logger.warning(
                "OllamaProvider: embed_model '%s' looks like an OpenAI/Voyage model "
                "which Ollama cannot use. Falling back to default '%s'.",
                candidate_embed,
                self.DEFAULT_EMBED_MODEL,
            )
            candidate_embed = self.DEFAULT_EMBED_MODEL
        self._embed_model = candidate_embed
        self._base_url = (
            base_url
            or os.environ.get("OLLAMA_BASE_URL", _DEFAULT_BASE_URL)
        ).rstrip("/")
        self._http_client: Any = None

    # --- internal helpers ---------------------------------------------------

    def _get_http(self) -> Any:
        """Return a shared aiohttp.ClientSession (lazy init)."""
        if self._http_client is None:
            try:
                import aiohttp  # type: ignore
            except ImportError:
                raise RuntimeError(
                    "aiohttp package not installed. Run: pip install aiohttp"
                )
            self._http_client = aiohttp.ClientSession()
        return self._http_client

    @staticmethod
    def _estimate_tokens(text: str) -> int:
        """Rough token count heuristic — 4 chars ≈ 1 token."""
        return max(1, len(text) // 4)

    # --- LLMProvider interface ----------------------------------------------

    async def complete(
        self,
        system: str,
        user: str,
        model: str | None = None,
        response_format: str = "text",
        max_tokens: int | None = None,
        temperature: float = 0.0,
    ) -> CompletionResponse:
        """
        Chat completion via Ollama /api/chat endpoint.

        If response_format == "json", sets format="json" in the request
        (Ollama will attempt to produce valid JSON but doesn't enforce schema).
        """
        resolved_model = self._resolve_model(model, self._completion_model)
        http = self._get_http()

        payload: dict[str, Any] = {
            "model": resolved_model,
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
            "stream": False,
            "options": {"temperature": temperature},
        }
        if response_format == "json":
            payload["format"] = "json"
        if max_tokens is not None:
            payload["options"]["num_predict"] = max_tokens

        url = f"{self._base_url}/api/chat"

        async def _do_chat():
            async with http.post(url, json=payload) as resp:
                resp.raise_for_status()
                return await resp.json()

        t0 = time.perf_counter()
        body = await retry_with_backoff(
            _do_chat,
            operation_name=f"ollama.complete({resolved_model})",
        )
        latency_ms = (time.perf_counter() - t0) * 1000

        text = body.get("message", {}).get("content", "")
        # Ollama returns eval_count (output) and prompt_eval_count (input)
        tokens_in = body.get("prompt_eval_count") or self._estimate_tokens(system + user)
        tokens_out = body.get("eval_count") or self._estimate_tokens(text)

        record = UsageRecord(
            provider=self.PROVIDER_NAME,
            model=resolved_model,
            call_type="complete",
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            cost_estimate=0.0,     # local — no cost
            latency_ms=latency_ms,
        )
        self._log_usage(record)

        return CompletionResponse(
            text=text,
            model=resolved_model,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            cost_estimate=0.0,
            latency_ms=latency_ms,
        )

    async def extract(
        self,
        system: str,
        user: str,
        schema: dict[str, Any],
        model: str | None = None,
        max_tokens: int | None = None,
        temperature: float = 0.0,
    ) -> ExtractionResponse:
        """
        Structured extraction via Ollama JSON mode.

        Ollama doesn't support tool-use / JSON schema enforcement, so we
        inject the schema as a prompt instruction and parse the output.
        Caller should validate the returned dict if schema compliance matters.
        """
        resolved_model = self._resolve_model(model, self._completion_model)

        schema_hint = json.dumps(schema, indent=2)
        augmented_system = (
            f"{system}\n\n"
            "Respond ONLY with a valid JSON object that matches this schema:\n"
            f"{schema_hint}\n"
            "Do not include any explanation or markdown fencing. Pure JSON only."
        )

        t0 = time.perf_counter()
        cr = await self.complete(
            system=augmented_system,
            user=user,
            model=resolved_model,
            response_format="json",
            max_tokens=max_tokens,
            temperature=temperature,
        )
        latency_ms = (time.perf_counter() - t0) * 1000

        try:
            data = json.loads(cr.text)
        except json.JSONDecodeError as exc:
            logger.warning(
                "OllamaProvider.extract: JSON parse failed — %s. Raw: %.200s",
                exc,
                cr.text,
            )
            data = {"_raw": cr.text, "_parse_error": str(exc)}

        record = UsageRecord(
            provider=self.PROVIDER_NAME,
            model=resolved_model,
            call_type="extract",
            tokens_in=cr.tokens_in,
            tokens_out=cr.tokens_out,
            cost_estimate=0.0,
            latency_ms=latency_ms,
        )
        self._log_usage(record)

        return ExtractionResponse(
            data=data,
            model=resolved_model,
            tokens_in=cr.tokens_in,
            tokens_out=cr.tokens_out,
            cost_estimate=0.0,
            latency_ms=latency_ms,
        )

    async def embed(
        self,
        texts: list[str],
        model: str | None = None,
        input_type: str = "query",   # ignored by Ollama, kept for interface compat
    ) -> EmbedResponse:
        """
        Embeddings via Ollama /api/embed endpoint.
        Processes texts individually (Ollama v0.3+ supports batch via list).
        """
        resolved_model = model or self._embed_model
        http = self._get_http()

        url = f"{self._base_url}/api/embed"
        payload = {
            "model": resolved_model,
            "input": texts,   # Ollama 0.3+ supports a list directly
        }

        async def _do_embed():
            async with http.post(url, json=payload) as resp:
                if resp.status == 400:
                    # Older Ollama: fall back to single-text mode
                    embs = []
                    for text in texts:
                        p = {"model": resolved_model, "prompt": text}
                        async with http.post(
                            f"{self._base_url}/api/embeddings", json=p
                        ) as r:
                            r.raise_for_status()
                            body = await r.json()
                            embs.append(body["embedding"])
                    return embs
                else:
                    resp.raise_for_status()
                    body = await resp.json()
                    return body.get("embeddings", [])

        t0 = time.perf_counter()
        embeddings = await retry_with_backoff(
            _do_embed,
            operation_name=f"ollama.embed({resolved_model})",
        )
        latency_ms = (time.perf_counter() - t0) * 1000

        tokens_in = sum(self._estimate_tokens(t) for t in texts)

        record = UsageRecord(
            provider=self.PROVIDER_NAME,
            model=resolved_model,
            call_type="embed",
            tokens_in=tokens_in,
            tokens_out=0,
            cost_estimate=0.0,
            latency_ms=latency_ms,
        )
        self._log_usage(record)

        return EmbedResponse(
            embeddings=embeddings,
            model=resolved_model,
            tokens_in=tokens_in,
            tokens_out=0,
            cost_estimate=0.0,
            latency_ms=latency_ms,
        )

    async def close(self) -> None:
        """Close the underlying HTTP session. Call when provider is no longer needed."""
        if self._http_client is not None:
            await self._http_client.close()
            self._http_client = None
