"""Ollama Adapter - Local model inference via Ollama.

Ollama provides a simple way to run LLMs locally:
- Install: https://ollama.ai
- Run: ollama serve
- Pull models: ollama pull llama3

Environment:
    OLLAMA_HOST: Base URL (default: http://localhost:11434)

Model naming:
    - "llama3", "mistral", "phi3", "qwen2" → local models
    - Models without "/" are assumed to be local
"""

from __future__ import annotations

import asyncio
import contextlib
import json
import logging
import os
from typing import Any

import aiohttp
from pydantic import BaseModel, Field

from core.types.llm import LLMEmbedResult, LLMGenerateResult, LLMToolCall, LLMToolCallFunction

logger = logging.getLogger(__name__)

# Models that should auto-route to Ollama (if Ollama is available)
OLLAMA_MODEL_PREFIXES = (
    "llama",
    "mistral",
    "phi",
    "qwen",
    "gemma",
    "codellama",
    "deepseek",
    "starcoder",
    "wizard",
    "vicuna",
    "orca",
    "neural",
    "dolphin",
    "nous",
    "tinyllama",
)


class OllamaModelInfo(BaseModel):
    """Info about a model available in Ollama."""

    model_config = {"extra": "allow"}

    name: str
    model: str = ""
    size: int = 0
    digest: str = ""
    modified_at: str = Field(default="", alias="modified_at")


class OllamaAdapter:
    """Adapter for Ollama local inference.

    Ollama exposes an OpenAI-compatible API at /v1/chat/completions,
    plus its native API at /api/generate and /api/chat.

    We use the native API for more control.
    """

    def __init__(
        self,
        host: str | None = None,
        timeout: float = 120.0,
    ):
        self.host = (host or os.getenv("OLLAMA_HOST", "http://localhost:11434")).rstrip("/")
        self.timeout = timeout
        self._session: aiohttp.ClientSession | None = None

    async def _get_session(self) -> aiohttp.ClientSession:
        # pytest / embedding runtimes may swap event loops between calls.
        # aiohttp sessions are bound to the loop they were created with.
        # If the loop changed (or is closed), recreate the session.
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = None

        if self._session is not None:
            try:
                sess_loop = getattr(self._session, "_loop", None)
                if (
                    sess_loop is not None
                    and hasattr(sess_loop, "is_closed")
                    and sess_loop.is_closed()
                ) or (loop is not None and sess_loop is not None and sess_loop is not loop):
                    await self.aclose()
            except Exception:
                # Best-effort cleanup; we'll recreate below if needed.
                with contextlib.suppress(Exception):
                    await self.aclose()

        if self._session is None or self._session.closed:
            self._session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=self.timeout))
        return self._session

    async def aclose(self) -> None:
        if self._session and not self._session.closed:
            await self._session.close()

    async def generate(
        self,
        *,
        prompt: str = "",
        model: str,
        system: str | None = None,
        temperature: float = 0.7,
        max_tokens: int = 4096,
        response_format: str = "text",
        tools: list[dict[str, Any]] | None = None,
        _tool_choice: str | dict[str, Any] | None = None,
        messages: list[dict[str, Any]] | None = None,
        images: list[str] | None = None,
    ) -> LLMGenerateResult:
        """Generate text using Ollama.

        Uses /api/chat for conversational interface.
        Messages, tools, and tool_choice arrive as dicts (the invoke.py layer
        converts Pydantic models before calling providers).
        """
        from core.llm.providers import LLMProviderError

        session = await self._get_session()

        # Messages are already dicts from invoke.py; use as-is.
        if messages:
            chat_messages = list(messages)
        else:
            chat_messages: list[dict[str, Any]] = []
            if system:
                chat_messages.append({"role": "system", "content": system})

            # Handle images for multimodal models (llava, etc.)
            if images:
                chat_messages.append(
                    {
                        "role": "user",
                        "content": prompt,
                        "images": images,  # Ollama accepts base64 images
                    },
                )
            else:
                chat_messages.append({"role": "user", "content": prompt})

        # Build request
        request_body: dict[str, Any] = {
            "model": model,
            "messages": chat_messages,
            "stream": False,
            "options": {
                "temperature": temperature,
                "num_predict": max_tokens,
            },
        }

        # JSON mode
        if response_format == "json":
            request_body["format"] = "json"

        # Tool calling (if supported by model)
        if tools:
            request_body["tools"] = self._convert_tools(tools)

        try:
            async with session.post(
                f"{self.host}/api/chat",
                json=request_body,
            ) as resp:
                if resp.status != 200:
                    error_text = await resp.text()
                    raise LLMProviderError(
                        message=f"Ollama error ({resp.status}): {error_text}",
                        provider="ollama",
                        model=model,
                        is_transient=resp.status >= 500,
                    )

                data = await resp.json()

        except aiohttp.ClientError as e:
            raise LLMProviderError(
                message=f"Connection error: {e}. Is Ollama running? (ollama serve)",
                provider="ollama",
                model=model,
                is_transient=True,
            ) from e

        # Parse response
        message = data.get("message", {})
        content = message.get("content", "")

        # Token counts (Ollama provides these)
        prompt_eval_count = data.get("prompt_eval_count", 0)
        eval_count = data.get("eval_count", 0)
        total_tokens = prompt_eval_count + eval_count

        # Tool calls (if any)
        tool_calls = None
        if message.get("tool_calls"):
            tool_calls = [
                LLMToolCall(
                    id=f"call_{i}",
                    function=LLMToolCallFunction(
                        name=tc.get("function", {}).get("name", ""),
                        arguments=json.dumps(tc.get("function", {}).get("arguments", {})),
                    ),
                )
                for i, tc in enumerate(message["tool_calls"])
            ]

        return LLMGenerateResult(
            content=content,
            model=model,
            tokens_used=total_tokens,
            cost_usd=0.0,  # Local models are free!
            input_tokens=prompt_eval_count,
            output_tokens=eval_count,
            provider="ollama",
            finish_reason=data.get("done_reason", "stop"),
            tool_calls=tool_calls,
            metadata={
                "total_duration": data.get("total_duration"),
                "load_duration": data.get("load_duration"),
                "eval_duration": data.get("eval_duration"),
            },
        )

    async def embed(
        self,
        texts: list[str],
        *,
        model: str = "nomic-embed-text",
    ) -> LLMEmbedResult:
        """Generate embeddings using Ollama.

        Requires an embedding model like nomic-embed-text.
        """
        from core.llm.providers import LLMProviderError

        session = await self._get_session()
        vectors = []
        total_tokens = 0

        for text in texts:
            request_body = {
                "model": model,
                "prompt": text,
            }

            try:
                async with session.post(
                    f"{self.host}/api/embeddings",
                    json=request_body,
                ) as resp:
                    if resp.status != 200:
                        error_text = await resp.text()
                        raise LLMProviderError(
                            message=f"Ollama embedding error: {error_text}",
                            provider="ollama",
                            model=model,
                            is_transient=resp.status >= 500,
                        )
                    data = await resp.json()
                    vectors.append(data.get("embedding", []))
            except aiohttp.ClientError as e:
                raise LLMProviderError(
                    message=f"Connection error: {e}",
                    provider="ollama",
                    model=model,
                    is_transient=True,
                ) from e

        return LLMEmbedResult(
            vectors=vectors,
            model=model,
            tokens_used=total_tokens,
            cost_usd=0.0,
            provider="ollama",
        )

    def _convert_tools(self, tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """Convert OpenAI-style tool dicts to Ollama wire format.

        Input shape (from invoke.py)::

            {"type": "function", "function": {"name": ..., "description": ..., "parameters": ...}}

        Ollama expects the same OpenAI-compatible shape, so pass through.
        """
        return tools

    async def list_models(self) -> list[OllamaModelInfo]:
        """List available models in Ollama."""
        session = await self._get_session()

        try:
            async with session.get(f"{self.host}/api/tags") as resp:
                if resp.status != 200:
                    return []
                data = await resp.json()
                return [OllamaModelInfo.model_validate(m) for m in data.get("models", [])]
        except Exception:
            return []

    async def is_available(self) -> bool:
        """Check if Ollama is running and accessible."""
        session = await self._get_session()

        try:
            async with session.get(
                f"{self.host}/api/tags",
                timeout=aiohttp.ClientTimeout(total=2.0),
            ) as resp:
                return resp.status == 200
        except Exception:
            return False


def is_ollama_model(model: str) -> bool:
    """Check if a model should be routed to Ollama."""
    model_lower = model.lower()
    # No slash = local model
    if "/" not in model_lower:
        return any(model_lower.startswith(prefix) for prefix in OLLAMA_MODEL_PREFIXES)
    return False


__all__ = [
    "OLLAMA_MODEL_PREFIXES",
    "OllamaAdapter",
    "OllamaModelInfo",
    "is_ollama_model",
]
