"""
Ollama local model provider.

Supports:
  - Any model available through Ollama (Llama 3, Mistral, Qwen, etc.)
  - Tool calling (for supported models)
  - Custom Ollama server endpoint
  - No API key required — fully local
"""

from __future__ import annotations

import json
import os
from typing import Any

import httpx

from nexus.core.message import Message
from nexus.models.base import ModelProvider
from nexus.types import TokenUsage


class OllamaProvider(ModelProvider):
    """
    LLM provider for Ollama local models.

    Usage:
        provider = OllamaProvider(model="llama3.1")
        response = await provider.chat(messages=[...])
    """

    def __init__(
        self,
        model: str = "llama3.1",
        base_url: str | None = None,
        default_temperature: float = 0.7,
        default_max_tokens: int | None = None,
        timeout: float = 120.0,
    ) -> None:
        self._model = model
        self._base_url = (
            base_url
            or os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434")
        ).rstrip("/")
        self._default_temperature = default_temperature
        self._default_max_tokens = default_max_tokens
        self._timeout = timeout
        self._client: httpx.AsyncClient | None = None

    def _get_client(self) -> httpx.AsyncClient:
        if self._client is None:
            self._client = httpx.AsyncClient(
                base_url=self._base_url,
                timeout=httpx.Timeout(self._timeout),
            )
        return self._client

    def _convert_tools(
        self, tools: list[dict[str, Any]]
    ) -> list[dict[str, Any]]:
        """
        Convert OpenAI-style tool definitions to Ollama format.
        Ollama uses the same format as OpenAI for function calling.
        """
        ollama_tools: list[dict[str, Any]] = []
        for tool in tools:
            if "function" in tool:
                ollama_tools.append({
                    "type": "function",
                    "function": tool["function"],
                })
            else:
                ollama_tools.append(tool)
        return ollama_tools

    async def chat(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        client = self._get_client()

        # Build Ollama /api/chat payload
        payload: dict[str, Any] = {
            "model": self._model,
            "messages": messages,
            "stream": False,
            "options": {
                "temperature": (
                    temperature
                    if temperature is not None
                    else self._default_temperature
                ),
            },
        }

        effective_max_tokens = max_tokens or self._default_max_tokens
        if effective_max_tokens:
            payload["options"]["num_predict"] = effective_max_tokens

        if tools:
            payload["tools"] = self._convert_tools(tools)

        # Merge extra kwargs into options
        if kwargs:
            payload["options"].update(kwargs)

        response = await client.post("/api/chat", json=payload)
        response.raise_for_status()
        data = response.json()

        resp_message = data.get("message", {})
        content = resp_message.get("content", "")

        # Parse tool calls from Ollama response
        tool_calls: list[dict[str, Any]] | None = None
        raw_tool_calls = resp_message.get("tool_calls")
        if raw_tool_calls:
            tool_calls = []
            for i, tc in enumerate(raw_tool_calls):
                func = tc.get("function", {})
                tool_calls.append({
                    "id": f"call_{i}",
                    "type": "function",
                    "function": {
                        "name": func.get("name", ""),
                        "arguments": (
                            json.dumps(func["arguments"])
                            if isinstance(func.get("arguments"), dict)
                            else func.get("arguments", "{}")
                        ),
                    },
                })

        # Build normalized Message
        message = Message.agent(
            content=content,
            agent_name=self._model,
        )
        if tool_calls:
            message.metadata["tool_calls"] = tool_calls

        # Token usage — Ollama provides these in response
        prompt_tokens = data.get("prompt_eval_count", 0) or 0
        completion_tokens = data.get("eval_count", 0) or 0
        usage = TokenUsage(
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            total_tokens=prompt_tokens + completion_tokens,
        )

        return {
            "message": message,
            "usage": usage,
            "model": data.get("model", self._model),
            "finish_reason": (
                "tool_calls" if tool_calls else data.get("done_reason", "stop")
            ),
            "tool_calls": tool_calls,
            "raw": data,
        }

    async def close(self) -> None:
        if self._client is not None:
            await self._client.aclose()
            self._client = None

    @property
    def model_name(self) -> str:
        return self._model
