"""Ollama LLM provider for local models."""

from __future__ import annotations

import json
from typing import Any

import httpx

from agenai.agent.llm import LLMProvider, LLMResponse, ToolCall


class OllamaProvider(LLMProvider):
    """Ollama provider for running local LLMs.

    Ollama is free and runs locally - no API key needed.
    Install: https://ollama.ai

    Supported models with tool use:
    - llama3.1 (recommended)
    - mistral
    - mixtral
    - qwen2.5

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

    def __init__(
        self,
        model: str = "llama3.1",
        base_url: str = "http://localhost:11434",
    ):
        """Initialize Ollama provider.

        Args:
            model: Model name (e.g., "llama3.1", "mistral", "qwen2.5")
            base_url: Ollama server URL (default: localhost:11434)
        """
        self.model = model
        self.base_url = base_url.rstrip("/")
        self.client = httpx.Client(timeout=120.0)

    def chat(
        self,
        messages: list[dict[str, Any]],
        system: str | None = None,
        tools: list[dict[str, Any]] | None = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
    ) -> LLMResponse:
        """Send chat completion request to Ollama."""

        # Build messages with system prompt
        ollama_messages = []

        if system:
            ollama_messages.append({"role": "system", "content": system})

        # Convert messages to Ollama format
        for msg in messages:
            ollama_messages.append(self._convert_message(msg))

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

        # Add tools if provided (Ollama supports tool use)
        if tools:
            request_body["tools"] = self._convert_tools(tools)

        # Make request
        response = self.client.post(
            f"{self.base_url}/api/chat",
            json=request_body,
        )
        response.raise_for_status()

        return self._parse_response(response.json())

    def _convert_message(self, msg: dict[str, Any]) -> dict[str, Any]:
        """Convert message to Ollama format."""
        role = msg.get("role", "user")
        content = msg.get("content", "")

        # Handle tool results
        if isinstance(content, list):
            # Check for tool_result type
            for item in content:
                if isinstance(item, dict) and item.get("type") == "tool_result":
                    return {
                        "role": "tool",
                        "content": item.get("content", ""),
                    }
            # Otherwise concatenate text
            text_parts = []
            for item in content:
                if isinstance(item, dict):
                    if item.get("type") == "text":
                        text_parts.append(item.get("text", ""))
                    elif item.get("type") == "tool_use":
                        # This is an assistant message with tool call
                        return {
                            "role": "assistant",
                            "content": "",
                            "tool_calls": [{
                                "function": {
                                    "name": item.get("name"),
                                    "arguments": json.dumps(item.get("input", {})),
                                }
                            }],
                        }
            content = "\n".join(text_parts)

        return {"role": role, "content": content}

    def _convert_tools(self, tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """Convert Anthropic tool format to Ollama format."""
        ollama_tools = []
        for tool in tools:
            ollama_tools.append({
                "type": "function",
                "function": {
                    "name": tool["name"],
                    "description": tool["description"],
                    "parameters": tool.get("input_schema", {}),
                },
            })
        return ollama_tools

    def _parse_response(self, response: dict[str, Any]) -> LLMResponse:
        """Parse Ollama response into LLMResponse."""
        message = response.get("message", {})
        content = message.get("content", "")
        tool_calls_raw = message.get("tool_calls", [])

        # Parse tool calls
        tool_calls = []
        for i, tc in enumerate(tool_calls_raw):
            func = tc.get("function", {})
            name = func.get("name", "")
            args_str = func.get("arguments", "{}")

            try:
                args = json.loads(args_str) if isinstance(args_str, str) else args_str
            except json.JSONDecodeError:
                args = {}

            tool_calls.append(ToolCall(
                id=f"call_{i}",
                name=name,
                arguments=args,
            ))

        return LLMResponse(
            text=content if content else None,
            tool_calls=tool_calls,
            stop_reason="tool_calls" if tool_calls else "stop",
            usage={
                "input_tokens": response.get("prompt_eval_count", 0),
                "output_tokens": response.get("eval_count", 0),
            },
        )

    def list_models(self) -> list[str]:
        """List available models in Ollama."""
        response = self.client.get(f"{self.base_url}/api/tags")
        response.raise_for_status()
        data = response.json()
        return [model["name"] for model in data.get("models", [])]

    def is_available(self) -> bool:
        """Check if Ollama server is running."""
        try:
            response = self.client.get(f"{self.base_url}/api/tags")
            return response.status_code == 200
        except Exception:
            return False
