"""
Ollama LLM client module
Provides local LLM inference using Ollama.
"""

import json
import logging
from typing import Any

import httpx

from .llm_interface import (
    ChatMessage,
    LLMFactory,
    LLMInterface,
    LLMResponse,
    MessageRole,
)


logger = logging.getLogger(__name__)


class OllamaClient(LLMInterface):
    """Ollama LLM client implementation for local inference"""

    def __init__(
        self,
        model: str = "llama3",
        base_url: str = "http://localhost:11434",
        timeout: int = 600,
        **kwargs: Any,
    ) -> None:
        """
        Initialize Ollama client
        Args:
            model: Model name (default: llama3)
            base_url: Ollama API base URL (default: http://localhost:11434)
            timeout: HTTP request timeout in seconds (default: 600 for large files)
            **kwargs: Other configuration parameters
        Raises:
            ConnectionError: If Ollama is not running
        """
        super().__init__(model=model, **kwargs)
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.client = httpx.AsyncClient(timeout=timeout)
        self._connection_checked = False

        logger.info(f"Ollama client created with model: {model}, timeout: {timeout}s")

    async def _check_connection(self) -> bool:
        """
        Check if Ollama is running and accessible
        Returns:
            True if connection successful, False otherwise
        """
        try:
            response = await self.client.get(f"{self.base_url}/api/tags")
            return response.status_code == 200
        except Exception as e:
            logger.error(f"Ollama connection check failed: {e}")
            return False

    async def _check_model(self) -> bool:
        """
        Check if the specified model is available
        Returns:
            True if model is available, False otherwise
        """
        try:
            response = await self.client.get(f"{self.base_url}/api/tags")
            if response.status_code == 200:
                data = response.json()
                models = data.get("models", [])
                # Get full model names (e.g., "glm-4.7-flash:q8_0")
                model_names = [m.get("name", "") for m in models]
                # Get base model names (e.g., "glm-4.7-flash")
                base_model_names = [name.split(":")[0] for name in model_names]

                # Check if exact match or base match
                if self.model in model_names:
                    return True
                # Check base name (e.g., "llama3" matches "llama3:latest")
                base_model = self.model.split(":")[0]
                if base_model in base_model_names:
                    return True

                logger.warning(
                    f"Model '{self.model}' not found in available models. "
                    f"Available: {', '.join(model_names[:5])}{'...' if len(model_names) > 5 else ''}"
                )
                return False
            return False
        except Exception as e:
            logger.warning(f"Model check failed: {e}")
            return True  # Assume available if check fails

    async def chat(self, messages: list[ChatMessage]) -> LLMResponse:
        """
        Chat interface
        Args:
            messages: List of chat messages
        Returns:
            LLM response
        Raises:
            Exception: If API call fails
        """
        try:
            # Convert messages to Ollama format
            ollama_messages = [
                {"role": msg.role.value, "content": msg.content} for msg in messages
            ]

            # Call Ollama chat API
            response = await self.client.post(
                f"{self.base_url}/api/chat",
                json={
                    "model": self.model,
                    "messages": ollama_messages,
                    "stream": False,
                },
            )

            response.raise_for_status()
            data = response.json()

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

            return LLMResponse(
                content=content,
                metadata={
                    "model": self.model,
                    "provider": "ollama",
                    "base_url": self.base_url,
                    "raw_response": data,
                },
                usage=None,
                model=self.model,
            )
        except Exception as e:
            logger.error(f"Ollama chat failed: {e}")
            raise

    async def complete(self, prompt: str, **kwargs: Any) -> LLMResponse:
        """
        Text completion interface
        Args:
            prompt: Input prompt
            **kwargs: Additional parameters
        Returns:
            LLM response
        Raises:
            Exception: If API call fails
        """
        try:
            # Use chat interface for completion
            messages = [ChatMessage(role=MessageRole.USER, content=prompt)]
            return await self.chat(messages)
        except Exception as e:
            logger.error(f"Ollama completion failed: {e}")
            raise

    async def extract_structured(
        self, content: str, schema: dict[str, Any], **kwargs: Any
    ) -> dict[str, Any]:
        """
        Structured data extraction interface
        Args:
            content: Input content
            schema: Output schema definition
            **kwargs: Additional parameters
        Returns:
            Extracted structured data
        Raises:
            Exception: If extraction fails
        """
        try:
            # Build extraction prompt
            prompt = self._build_extraction_prompt(content, schema)

            # Call completion interface
            response = await self.complete(prompt, **kwargs)

            # Try to parse as JSON
            response_text = response.content.strip()

            # Find JSON in response (handle potential extra text)
            if response_text.startswith("```json"):
                response_text = response_text[7:]
            if response_text.startswith("```"):
                response_text = response_text[3:]
            if response_text.endswith("```"):
                response_text = response_text[:-3]
            response_text = response_text.strip()

            try:
                extracted_data = json.loads(response_text)
                return {
                    "extracted_data": extracted_data,
                    "schema": schema,
                    "model": self.model,
                    "provider": "ollama",
                }
            except json.JSONDecodeError:
                # If not valid JSON, return raw content
                return {
                    "extracted_data": response_text,
                    "schema": schema,
                    "model": self.model,
                    "provider": "ollama",
                    "parse_error": "Could not parse as JSON",
                }
        except Exception as e:
            logger.error(f"Ollama structured extraction failed: {e}")
            raise

    def _build_extraction_prompt(self, content: str, schema: dict[str, Any]) -> str:
        """Build extraction prompt for structured data"""
        prompt = f"""Extract structured information from the following content according to the schema.

Content:
{content}

Schema:
{schema}

Please return the extracted data in valid JSON format. If a field is not found, use null.
Return ONLY the JSON object without any additional text or formatting."""
        return prompt

    def get_model_info(self) -> dict[str, Any]:
        """Get model information"""
        info = super().get_model_info()
        info.update(
            {
                "provider": "Ollama",
                "base_url": self.base_url,
                "local_inference": True,
            }
        )
        return info

    async def test_connection(self) -> bool:
        """
        Test Ollama connection
        Returns:
            True if connection successful, False otherwise
        """
        return await self._check_connection()

    async def __aenter__(self):
        """Async context manager entry"""
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async context manager exit"""
        _ = exc_type, exc_val, exc_tb
        await self.client.aclose()


class OllamaStreamClient(OllamaClient):
    """Ollama streaming client for real-time responses"""

    async def chat_stream(self, messages: list[ChatMessage]) -> Any:
        """
        Streaming chat interface
        Args:
            messages: List of chat messages
        Yields:
            Streaming response chunks
        Raises:
            Exception: If streaming fails
        """
        try:
            # Convert messages to Ollama format
            ollama_messages = [
                {"role": msg.role.value, "content": msg.content} for msg in messages
            ]

            # Stream responses from Ollama
            async with self.client.stream(
                "POST",
                f"{self.base_url}/api/chat",
                json={
                    "model": self.model,
                    "messages": ollama_messages,
                    "stream": True,
                },
                timeout=None,
            ) as response:
                async for line in response.aiter_lines():
                    if line.strip():
                        try:
                            chunk = json.loads(line)
                            content = chunk.get("message", {}).get("content", "")
                            if content:
                                yield LLMResponse(
                                    content=content,
                                    metadata={
                                        "model": self.model,
                                        "chunk_type": "delta",
                                        "provider": "ollama",
                                    },
                                    model=self.model,
                                )
                        except json.JSONDecodeError:
                            continue
        except Exception as e:
            logger.error(f"Ollama streaming chat failed: {e}")
            raise


# Register Ollama providers
LLMFactory.register("ollama", OllamaClient)
LLMFactory.register("ollama_stream", OllamaStreamClient)


# Convenience functions
def create_ollama_client(
    model: str = "llama3", base_url: str = "http://localhost:11434", **kwargs: Any
) -> OllamaClient:
    """
    Create Ollama client
    Args:
        model: Model name
        base_url: Ollama API base URL
        **kwargs: Configuration parameters
    Returns:
        Ollama client instance
    """
    return OllamaClient(model=model, base_url=base_url, **kwargs)


def create_ollama_stream_client(
    model: str = "llama3", base_url: str = "http://localhost:11434", **kwargs: Any
) -> OllamaStreamClient:
    """
    Create Ollama streaming client
    Args:
        model: Model name
        base_url: Ollama API base URL
        **kwargs: Configuration parameters
    Returns:
        Ollama streaming client instance
    """
    return OllamaStreamClient(model=model, base_url=base_url, **kwargs)
