"""Ollama provider implementation for Oclawma."""

from __future__ import annotations

import json
from collections.abc import AsyncIterator
from typing import Any

import httpx

from oclawma.providers.base import (
    AuthenticationError,
    BaseProvider,
    CompletionError,
    CompletionRequest,
    CompletionResponse,
    Message,
    ModelNotFoundError,
    ProviderConnectionError,
    ProviderError,
    RateLimitError,
    UsageStats,
)


class OllamaProvider(BaseProvider):
    """Ollama provider for local LLM inference."""

    DEFAULT_BASE_URL = "http://localhost:11434"
    # Approximate characters per token for estimation
    CHARS_PER_TOKEN = 4

    def __init__(
        self, base_url: str | None = None, api_key: str | None = None, timeout: float = 60.0
    ) -> None:
        """Initialize the Ollama provider.

        Args:
            base_url: The base URL for Ollama (default: http://localhost:11434).
            api_key: Optional API key (not typically needed for local Ollama).
            timeout: Request timeout in seconds (default: 60).
        """
        super().__init__(base_url or self.DEFAULT_BASE_URL, api_key)
        self._timeout = timeout
        self._client: httpx.AsyncClient | None = None

    async def _get_client(self) -> httpx.AsyncClient:
        """Get or create the HTTP client."""
        if self._client is None or self._client.is_closed:
            self._client = httpx.AsyncClient(
                base_url=self.base_url,
                timeout=httpx.Timeout(self._timeout, connect=10.0),
                headers={"Content-Type": "application/json"},
            )
        return self._client

    async def close(self) -> None:
        """Close the HTTP client connection."""
        if self._client and not self._client.is_closed:
            await self._client.aclose()
            self._client = None

    def _convert_messages(self, messages: list[Message]) -> tuple[str | None, list[dict[str, Any]]]:
        """Convert messages to Ollama format.

        Ollama uses a system prompt plus message history format.

        Returns:
            Tuple of (system_prompt, ollama_messages).
        """
        system_prompt = None
        ollama_messages = []

        for msg in messages:
            if msg.role == "system":
                system_prompt = msg.content
            else:
                ollama_messages.append(
                    {
                        "role": msg.role,
                        "content": msg.content,
                    }
                )

        return system_prompt, ollama_messages

    def _handle_error(self, error: httpx.HTTPError) -> ProviderError:
        """Convert HTTP errors to provider errors."""
        if isinstance(error, httpx.ConnectError):
            return ProviderConnectionError(
                f"Cannot connect to Ollama at {self.base_url}. "
                "Ensure Ollama is running and accessible.",
                cause=error,
            )
        elif isinstance(error, httpx.TimeoutException):
            return ProviderConnectionError(
                "Request to Ollama timed out. The model may be loading or the request too complex.",
                cause=error,
            )
        elif isinstance(error, httpx.HTTPStatusError):
            status_code = error.response.status_code
            if status_code == 401:
                return AuthenticationError(
                    "Authentication failed. Check your API key.", cause=error
                )
            elif status_code == 404:
                return ModelNotFoundError(
                    "Model not found. Ensure the model is pulled: ollama pull <model>",
                    cause=error,
                )
            elif status_code == 429:
                return RateLimitError("Rate limit exceeded. Please retry later.", cause=error)
            else:
                return CompletionError(
                    f"Ollama returned HTTP {status_code}: {error.response.text}",
                    cause=error,
                )
        else:
            return CompletionError(f"Unexpected error: {error}", cause=error)

    def count_tokens(self, text: str, model: str | None = None) -> int:
        """Estimate token count using character-based heuristic.

        Ollama doesn't expose tokenizers directly, so we use the
        approximation of ~4 characters per token for English text.

        Args:
            text: The text to count tokens for.
            model: Ignored for Ollama (no per-model tokenizers available).

        Returns:
            Estimated token count.
        """
        if not text:
            return 0
        # Use ceiling to ensure we don't underestimate
        return max(1, (len(text) + self.CHARS_PER_TOKEN - 1) // self.CHARS_PER_TOKEN)

    async def complete(self, request: CompletionRequest) -> CompletionResponse:
        """Generate a chat completion.

        Args:
            request: The completion request.

        Returns:
            The completion response.

        Raises:
            ConnectionError: If unable to connect to Ollama.
            ModelNotFoundError: If the model is not available.
            CompletionError: If completion generation fails.
        """
        client = await self._get_client()
        system_prompt, messages = self._convert_messages(request.messages)

        payload: dict[str, Any] = {
            "model": request.model,
            "messages": messages,
            "stream": False,
            "options": {
                "temperature": request.temperature,
                "top_p": request.top_p,
            },
        }

        if system_prompt:
            payload["system"] = system_prompt

        if request.max_tokens:
            payload["options"]["num_predict"] = request.max_tokens

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

            content = data.get("message", {}).get("content", "")
            # Ollama doesn't return token counts, so we estimate
            prompt_tokens = self.count_message_tokens(request.messages, request.model)
            completion_tokens = self.count_tokens(content, request.model)

            return CompletionResponse(
                content=content,
                model=request.model,
                usage=UsageStats(
                    prompt_tokens=prompt_tokens,
                    completion_tokens=completion_tokens,
                    total_tokens=prompt_tokens + completion_tokens,
                ),
                finish_reason="stop" if data.get("done", False) else None,
            )

        except httpx.HTTPError as e:
            raise self._handle_error(e) from e

    async def stream_complete(
        self, request: CompletionRequest
    ) -> AsyncIterator[CompletionResponse]:
        """Generate a streaming chat completion.

        Args:
            request: The completion request.

        Yields:
            Completion response chunks.

        Raises:
            ConnectionError: If unable to connect to Ollama.
            ModelNotFoundError: If the model is not available.
            CompletionError: If completion generation fails.
        """
        client = await self._get_client()
        system_prompt, messages = self._convert_messages(request.messages)

        payload: dict[str, Any] = {
            "model": request.model,
            "messages": messages,
            "stream": True,
            "options": {
                "temperature": request.temperature,
                "top_p": request.top_p,
            },
        }

        if system_prompt:
            payload["system"] = system_prompt

        if request.max_tokens:
            payload["options"]["num_predict"] = request.max_tokens

        prompt_tokens = self.count_message_tokens(request.messages, request.model)
        accumulated_content = ""

        try:
            async with client.stream("POST", "/api/chat", json=payload) as response:
                response.raise_for_status()

                async for line in response.aiter_lines():
                    if not line.strip():
                        continue

                    try:
                        chunk = json.loads(line)
                    except json.JSONDecodeError as e:
                        raise CompletionError(f"Invalid JSON in stream: {line}") from e

                    message = chunk.get("message", {})
                    content_chunk = message.get("content", "")
                    accumulated_content += content_chunk

                    # Only yield non-empty content chunks
                    if content_chunk:
                        completion_tokens = self.count_tokens(accumulated_content, request.model)
                        yield CompletionResponse(
                            content=content_chunk,
                            model=request.model,
                            usage=UsageStats(
                                prompt_tokens=prompt_tokens,
                                completion_tokens=completion_tokens,
                                total_tokens=prompt_tokens + completion_tokens,
                            ),
                            finish_reason=None,  # Will be set on final chunk
                        )

                    if chunk.get("done", False):
                        # Final chunk
                        completion_tokens = self.count_tokens(accumulated_content, request.model)
                        yield CompletionResponse(
                            content="",  # Empty for final chunk
                            model=request.model,
                            usage=UsageStats(
                                prompt_tokens=prompt_tokens,
                                completion_tokens=completion_tokens,
                                total_tokens=prompt_tokens + completion_tokens,
                            ),
                            finish_reason="stop",
                        )
                        break

        except httpx.HTTPError as e:
            raise self._handle_error(e) from e

    async def list_models(self) -> list[str]:
        """List available models from Ollama.

        Returns:
            List of available model names.

        Raises:
            ConnectionError: If unable to connect to Ollama.
        """
        client = await self._get_client()

        try:
            response = await client.get("/api/tags")
            response.raise_for_status()
            data = response.json()

            models = []
            for model in data.get("models", []):
                name = model.get("name", "")
                if name:
                    models.append(name)

            return sorted(models)

        except httpx.HTTPError as e:
            raise self._handle_error(e) from e

    async def health_check(self) -> dict[str, Any]:
        """Check the health of the Ollama connection.

        Returns:
            Health status information including version and available models.
        """
        result: dict[str, Any] = {
            "status": "unknown",
            "base_url": self.base_url,
            "version": None,
            "models": [],
            "error": None,
        }

        try:
            client = await self._get_client()

            # Try to get version
            try:
                response = await client.get("/api/version")
                if response.status_code == 200:
                    version_data = response.json()
                    result["version"] = version_data.get("version")
            except httpx.HTTPError:
                pass  # Version endpoint may not exist in older versions

            # List models to verify connection
            models = await self.list_models()
            result["models"] = models
            result["status"] = "healthy"

        except ProviderConnectionError as e:
            result["status"] = "unhealthy"
            result["error"] = str(e.message)
        except Exception as e:
            result["status"] = "error"
            result["error"] = str(e)

        return result

    async def pull_model(self, model: str) -> AsyncIterator[dict[str, Any]]:
        """Pull a model from Ollama.

        Args:
            model: The model name to pull.

        Yields:
            Progress updates during the pull.

        Raises:
            ConnectionError: If unable to connect to Ollama.
            CompletionError: If the pull fails.
        """
        client = await self._get_client()

        payload = {"name": model, "stream": True}

        try:
            async with client.stream("POST", "/api/pull", json=payload) as response:
                response.raise_for_status()

                async for line in response.aiter_lines():
                    if not line.strip():
                        continue

                    try:
                        chunk = json.loads(line)
                        yield chunk
                    except json.JSONDecodeError:
                        yield {"status": "unknown", "raw": line}

        except httpx.HTTPError as e:
            raise self._handle_error(e) from e
