"""
LLM Client - Unified interface for local and cloud LLM providers
Supports: Ollama (local GPU), Claude API, OpenAI API
"""

from abc import ABC, abstractmethod
from typing import AsyncGenerator, Dict, List, Optional, Any
from dataclasses import dataclass
import aiohttp
import json
from loguru import logger


@dataclass
class LLMMessage:
    """Standard message format"""
    role: str  # 'system', 'user', 'assistant'
    content: str
    tool_calls: Optional[List[Dict]] = None


@dataclass
class LLMResponse:
    """Standard response format"""
    content: str
    finish_reason: str
    tool_calls: Optional[List[Dict]] = None
    usage: Optional[Dict[str, int]] = None


class BaseLLMClient(ABC):
    """Abstract base class for LLM clients"""

    @abstractmethod
    async def generate(
        self,
        messages: List[LLMMessage],
        tools: Optional[List[Dict]] = None,
        temperature: float = 0.7,
        max_tokens: int = 2048,
        stream: bool = False
    ) -> LLMResponse | AsyncGenerator[str, None]:
        """Generate completion from messages"""
        pass

    @abstractmethod
    async def health_check(self) -> bool:
        """Check if LLM service is available"""
        pass


class OllamaClient(BaseLLMClient):
    """
    Ollama client for local LLM inference
    Supports: Llama 3, Mistral, and other Ollama models
    """

    def __init__(
        self,
        base_url: str = "http://localhost:11434",
        model: str = "llama3:8b",
        timeout: int = 300
    ):
        self.base_url = base_url
        self.model = model
        self.timeout = aiohttp.ClientTimeout(total=timeout)
        logger.info(f"Initialized Ollama client: {model} @ {base_url}")

    async def health_check(self) -> bool:
        """Check if Ollama is running"""
        try:
            async with aiohttp.ClientSession(timeout=self.timeout) as session:
                async with session.get(f"{self.base_url}/api/tags") as response:
                    return response.status == 200
        except Exception as e:
            logger.error(f"Ollama health check failed: {e}")
            return False

    async def generate(
        self,
        messages: List[LLMMessage],
        tools: Optional[List[Dict]] = None,
        temperature: float = 0.7,
        max_tokens: int = 2048,
        stream: bool = False
    ) -> LLMResponse | AsyncGenerator[str, None]:
        """
        Generate completion using Ollama

        Args:
            messages: List of conversation messages
            tools: Optional list of function calling tools
            temperature: Sampling temperature (0.0-2.0)
            max_tokens: Maximum tokens to generate
            stream: Whether to stream the response
        """

        # Convert messages to Ollama format
        ollama_messages = self._convert_messages(messages)

        # Build request payload
        payload = {
            "model": self.model,
            "messages": ollama_messages,
            "stream": stream,
            "options": {
                "temperature": temperature,
                "num_predict": max_tokens,
            }
        }

        # Add tools if provided (for function calling)
        if tools:
            payload["tools"] = self._convert_tools(tools)

        if stream:
            return self._generate_stream(payload)
        else:
            return await self._generate_complete(payload)

    async def _generate_complete(self, payload: Dict) -> LLMResponse:
        """Generate complete response (non-streaming)"""
        async with aiohttp.ClientSession(timeout=self.timeout) as session:
            async with session.post(
                f"{self.base_url}/api/chat",
                json=payload
            ) as response:
                if response.status != 200:
                    error_text = await response.text()
                    raise Exception(f"Ollama API error: {error_text}")

                result = await response.json()

                # Extract message content
                message = result.get("message", {})
                content = message.get("content", "")
                tool_calls = message.get("tool_calls")

                return LLMResponse(
                    content=content,
                    finish_reason=result.get("done_reason", "stop"),
                    tool_calls=tool_calls,
                    usage={
                        "prompt_tokens": result.get("prompt_eval_count", 0),
                        "completion_tokens": result.get("eval_count", 0),
                        "total_tokens": result.get("prompt_eval_count", 0) + result.get("eval_count", 0)
                    }
                )

    async def _generate_stream(self, payload: Dict) -> AsyncGenerator[str, None]:
        """Generate streaming response"""
        async with aiohttp.ClientSession(timeout=self.timeout) as session:
            async with session.post(
                f"{self.base_url}/api/chat",
                json=payload
            ) as response:
                if response.status != 200:
                    error_text = await response.text()
                    raise Exception(f"Ollama API error: {error_text}")

                async for line in response.content:
                    if line:
                        try:
                            chunk = json.loads(line)
                            message = chunk.get("message", {})
                            content = message.get("content", "")

                            if content:
                                yield content

                            # Check if done
                            if chunk.get("done", False):
                                break
                        except json.JSONDecodeError:
                            continue

    def _convert_messages(self, messages: List[LLMMessage]) -> List[Dict]:
        """Convert standard messages to Ollama format"""
        return [
            {
                "role": msg.role,
                "content": msg.content,
                **({"tool_calls": msg.tool_calls} if msg.tool_calls else {})
            }
            for msg in messages
        ]

    def _convert_tools(self, tools: List[Dict]) -> List[Dict]:
        """Convert OpenAI-style tools to Ollama format"""
        # Ollama uses similar format to OpenAI function calling
        return [
            {
                "type": "function",
                "function": tool.get("function", tool)
            }
            for tool in tools
        ]

    async def list_models(self) -> List[str]:
        """List available Ollama models"""
        async with aiohttp.ClientSession(timeout=self.timeout) as session:
            async with session.get(f"{self.base_url}/api/tags") as response:
                if response.status == 200:
                    data = await response.json()
                    return [model["name"] for model in data.get("models", [])]
                return []

    async def pull_model(self, model_name: str) -> bool:
        """Pull a model from Ollama library"""
        logger.info(f"Pulling model: {model_name}")

        async with aiohttp.ClientSession(timeout=self.timeout) as session:
            async with session.post(
                f"{self.base_url}/api/pull",
                json={"name": model_name, "stream": False}
            ) as response:
                if response.status == 200:
                    logger.info(f"Model {model_name} pulled successfully")
                    return True
                else:
                    error = await response.text()
                    logger.error(f"Failed to pull model: {error}")
                    return False


class ClaudeClient(BaseLLMClient):
    """
    Claude API client (for comparison/fallback)
    """

    def __init__(self, api_key: str, model: str = "claude-3-5-sonnet-20241022"):
        self.api_key = api_key
        self.model = model
        logger.info(f"Initialized Claude client: {model}")

    async def health_check(self) -> bool:
        """Check if Claude API is accessible"""
        # Simple validation
        return bool(self.api_key)

    async def generate(
        self,
        messages: List[LLMMessage],
        tools: Optional[List[Dict]] = None,
        temperature: float = 0.7,
        max_tokens: int = 2048,
        stream: bool = False
    ) -> LLMResponse | AsyncGenerator[str, None]:
        """Generate using Claude API"""
        # This would use the Anthropic SDK
        # Placeholder for now
        raise NotImplementedError("Claude client to be implemented with Anthropic SDK")


class LLMClientFactory:
    """Factory for creating LLM clients"""

    @staticmethod
    def create_client(
        provider: str = "ollama",
        **kwargs
    ) -> BaseLLMClient:
        """
        Create LLM client based on provider

        Args:
            provider: 'ollama', 'claude', or 'openai'
            **kwargs: Provider-specific configuration

        Returns:
            Configured LLM client instance
        """

        if provider.lower() == "ollama":
            return OllamaClient(
                base_url=kwargs.get("base_url", "http://localhost:11434"),
                model=kwargs.get("model", "llama3:8b"),
                timeout=kwargs.get("timeout", 300)
            )

        elif provider.lower() == "claude":
            return ClaudeClient(
                api_key=kwargs.get("api_key"),
                model=kwargs.get("model", "claude-3-5-sonnet-20241022")
            )

        else:
            raise ValueError(f"Unsupported LLM provider: {provider}")
