"""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
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,
    ):
        if host is None:
            from core.config import get_settings

            host = get_settings().ollama_host
        self.host = host.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)
            if tool_choice is not None:
                logger.warning(
                    "Ollama does not support tool_choice (%r); the parameter will be ignored. "
                    "Model: %s",
                    tool_choice,
                    model,
                )

        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 generate_stream(
        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,
    ):
        """Stream text generation via Ollama's native streaming API.

        Yields
        ------
        LLMStreamChunk
            Incremental chunks with ``delta`` text.
        """
        from core.llm.providers import LLMProviderError
        from core.types.llm import LLMStreamChunk

        session = await self._get_session()

        # Build messages (same logic as generate)
        if messages:
            chat_messages = list(messages)
        else:
            chat_messages: list[dict[str, Any]] = []
            if system:
                chat_messages.append({"role": "system", "content": system})
            if images:
                chat_messages.append(
                    {"role": "user", "content": prompt, "images": images},
                )
            else:
                chat_messages.append({"role": "user", "content": prompt})

        request_body: dict[str, Any] = {
            "model": model,
            "messages": chat_messages,
            "stream": True,
            "options": {
                "temperature": temperature,
                "num_predict": max_tokens,
            },
        }
        if response_format == "json":
            request_body["format"] = "json"
        if tools:
            request_body["tools"] = self._convert_tools(tools)
            if tool_choice is not None:
                logger.warning(
                    "Ollama does not support tool_choice (%r); the parameter will be ignored. "
                    "Model: %s",
                    tool_choice,
                    model,
                )

        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 streaming error ({resp.status}): {error_text}",
                        provider="ollama",
                        model=model,
                        is_transient=resp.status >= 500,
                    )

                async for raw_line in resp.content:
                    line = raw_line.decode("utf-8").strip()
                    if not line:
                        continue
                    data = json.loads(line)
                    msg = data.get("message", {})

                    # Tool calls (Ollama sends them in the final message)
                    tc_list = None
                    if msg.get("tool_calls"):
                        tc_list = [
                            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(msg["tool_calls"])
                        ]

                    done_reason = data.get("done_reason")
                    yield LLMStreamChunk(
                        delta=msg.get("content", "") or "",
                        finish_reason=done_reason if data.get("done") else None,
                        tool_calls=tc_list,
                    )

        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

    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",
]
