"""
Zubbl LLM Adapter - Model-Agnostic LLM Interface

Inspired by EvoAgentX's model abstraction layer.
Provides a unified interface for any LLM provider.

Attribution: Architecture patterns inspired by EvoAgentX (MIT License)
https://github.com/EvoAgentX/EvoAgentX
"""

import os
from abc import ABC, abstractmethod
from enum import Enum
from typing import Any, Dict, List, Optional, Union

from pydantic import BaseModel, Field


class LLMProvider(str, Enum):
    """Supported LLM providers"""

    OPENAI = "openai"
    ANTHROPIC = "anthropic"
    GOOGLE = "google"
    GROQ = "groq"
    OLLAMA = "ollama"
    TOGETHER = "together"
    OPENROUTER = "openrouter"
    CUSTOM = "custom"


class Message(BaseModel):
    """Chat message format"""

    role: str  # "system", "user", "assistant"
    content: str
    name: Optional[str] = None


class LLMConfig(BaseModel):
    """LLM configuration"""

    provider: LLMProvider = LLMProvider.OPENAI
    model: str = "gpt-4o-mini"
    api_key: Optional[str] = None
    base_url: Optional[str] = None
    temperature: float = 0.7
    max_tokens: int = 4096
    top_p: float = 1.0
    timeout: float = 60.0


class LLMResponse(BaseModel):
    """Standardized LLM response"""

    content: str
    model: str
    provider: LLMProvider
    usage: Dict[str, int] = Field(default_factory=dict)
    latency_ms: int = 0
    raw_response: Optional[Dict[str, Any]] = None


class BaseLLM(ABC):
    """Base class for all LLM implementations"""

    def __init__(self, config: LLMConfig):
        self.config = config

    @abstractmethod
    def generate(self, messages: List[Message], **kwargs) -> LLMResponse:
        """Generate a response from the LLM"""
        pass

    @abstractmethod
    async def generate_async(self, messages: List[Message], **kwargs) -> LLMResponse:
        """Async version of generate"""
        pass

    def chat(self, prompt: str, system: Optional[str] = None, **kwargs) -> str:
        """Simple chat interface - returns just the content"""
        messages = []
        if system:
            messages.append(Message(role="system", content=system))
        messages.append(Message(role="user", content=prompt))
        response = self.generate(messages, **kwargs)
        return response.content


class OpenAILLM(BaseLLM):
    """OpenAI LLM implementation"""

    def __init__(self, config: LLMConfig):
        super().__init__(config)
        self._client = None
        self._async_client = None

    def _get_client(self):
        if self._client is None:
            try:
                from openai import OpenAI

                self._client = OpenAI(
                    api_key=self.config.api_key or os.getenv("OPENAI_API_KEY"),
                    base_url=self.config.base_url,
                    timeout=self.config.timeout,
                )
            except ImportError:
                raise ImportError("openai package required. Install with: pip install openai")
        return self._client

    def _get_async_client(self):
        if self._async_client is None:
            try:
                from openai import AsyncOpenAI

                self._async_client = AsyncOpenAI(
                    api_key=self.config.api_key or os.getenv("OPENAI_API_KEY"),
                    base_url=self.config.base_url,
                    timeout=self.config.timeout,
                )
            except ImportError:
                raise ImportError("openai package required. Install with: pip install openai")
        return self._async_client

    def generate(self, messages: List[Message], **kwargs) -> LLMResponse:
        import time

        start = time.time()

        client = self._get_client()
        response = client.chat.completions.create(
            model=self.config.model,
            messages=[{"role": m.role, "content": m.content} for m in messages],
            temperature=kwargs.get("temperature", self.config.temperature),
            max_tokens=kwargs.get("max_tokens", self.config.max_tokens),
            top_p=kwargs.get("top_p", self.config.top_p),
        )

        latency = int((time.time() - start) * 1000)

        return LLMResponse(
            content=response.choices[0].message.content,
            model=self.config.model,
            provider=LLMProvider.OPENAI,
            usage={
                "prompt_tokens": response.usage.prompt_tokens,
                "completion_tokens": response.usage.completion_tokens,
                "total_tokens": response.usage.total_tokens,
            },
            latency_ms=latency,
            raw_response=response.model_dump(),
        )

    async def generate_async(self, messages: List[Message], **kwargs) -> LLMResponse:
        import time

        start = time.time()

        client = self._get_async_client()
        response = await client.chat.completions.create(
            model=self.config.model,
            messages=[{"role": m.role, "content": m.content} for m in messages],
            temperature=kwargs.get("temperature", self.config.temperature),
            max_tokens=kwargs.get("max_tokens", self.config.max_tokens),
            top_p=kwargs.get("top_p", self.config.top_p),
        )

        latency = int((time.time() - start) * 1000)

        return LLMResponse(
            content=response.choices[0].message.content,
            model=self.config.model,
            provider=LLMProvider.OPENAI,
            usage={
                "prompt_tokens": response.usage.prompt_tokens,
                "completion_tokens": response.usage.completion_tokens,
                "total_tokens": response.usage.total_tokens,
            },
            latency_ms=latency,
            raw_response=response.model_dump(),
        )


class AnthropicLLM(BaseLLM):
    """Anthropic Claude LLM implementation"""

    def __init__(self, config: LLMConfig):
        super().__init__(config)
        self._client = None
        self._async_client = None

    def _get_client(self):
        if self._client is None:
            try:
                from anthropic import Anthropic

                self._client = Anthropic(
                    api_key=self.config.api_key or os.getenv("ANTHROPIC_API_KEY"),
                )
            except ImportError:
                raise ImportError("anthropic package required. Install with: pip install anthropic")
        return self._client

    def _get_async_client(self):
        if self._async_client is None:
            try:
                from anthropic import AsyncAnthropic

                self._async_client = AsyncAnthropic(
                    api_key=self.config.api_key or os.getenv("ANTHROPIC_API_KEY"),
                )
            except ImportError:
                raise ImportError("anthropic package required. Install with: pip install anthropic")
        return self._async_client

    def generate(self, messages: List[Message], **kwargs) -> LLMResponse:
        import time

        start = time.time()

        client = self._get_client()

        # Extract system message if present
        system = None
        chat_messages = []
        for m in messages:
            if m.role == "system":
                system = m.content
            else:
                chat_messages.append({"role": m.role, "content": m.content})

        response = client.messages.create(
            model=self.config.model,
            messages=chat_messages,
            system=system,
            temperature=kwargs.get("temperature", self.config.temperature),
            max_tokens=kwargs.get("max_tokens", self.config.max_tokens),
        )

        latency = int((time.time() - start) * 1000)

        return LLMResponse(
            content=response.content[0].text,
            model=self.config.model,
            provider=LLMProvider.ANTHROPIC,
            usage={
                "prompt_tokens": response.usage.input_tokens,
                "completion_tokens": response.usage.output_tokens,
                "total_tokens": response.usage.input_tokens + response.usage.output_tokens,
            },
            latency_ms=latency,
            raw_response=response.model_dump(),
        )

    async def generate_async(self, messages: List[Message], **kwargs) -> LLMResponse:
        import time

        start = time.time()

        client = self._get_async_client()

        system = None
        chat_messages = []
        for m in messages:
            if m.role == "system":
                system = m.content
            else:
                chat_messages.append({"role": m.role, "content": m.content})

        response = await client.messages.create(
            model=self.config.model,
            messages=chat_messages,
            system=system,
            temperature=kwargs.get("temperature", self.config.temperature),
            max_tokens=kwargs.get("max_tokens", self.config.max_tokens),
        )

        latency = int((time.time() - start) * 1000)

        return LLMResponse(
            content=response.content[0].text,
            model=self.config.model,
            provider=LLMProvider.ANTHROPIC,
            usage={
                "prompt_tokens": response.usage.input_tokens,
                "completion_tokens": response.usage.output_tokens,
                "total_tokens": response.usage.input_tokens + response.usage.output_tokens,
            },
            latency_ms=latency,
            raw_response=response.model_dump(),
        )


class OllamaLLM(BaseLLM):
    """Ollama local LLM implementation"""

    def __init__(self, config: LLMConfig):
        super().__init__(config)
        self.base_url = config.base_url or "http://localhost:11434"

    def generate(self, messages: List[Message], **kwargs) -> LLMResponse:
        import time

        import httpx

        start = time.time()

        response = httpx.post(
            f"{self.base_url}/api/chat",
            json={
                "model": self.config.model,
                "messages": [{"role": m.role, "content": m.content} for m in messages],
                "stream": False,
                "options": {
                    "temperature": kwargs.get("temperature", self.config.temperature),
                    "num_predict": kwargs.get("max_tokens", self.config.max_tokens),
                },
            },
            timeout=self.config.timeout,
        )
        response.raise_for_status()
        data = response.json()

        latency = int((time.time() - start) * 1000)

        return LLMResponse(
            content=data["message"]["content"],
            model=self.config.model,
            provider=LLMProvider.OLLAMA,
            usage={
                "prompt_tokens": data.get("prompt_eval_count", 0),
                "completion_tokens": data.get("eval_count", 0),
                "total_tokens": data.get("prompt_eval_count", 0) + data.get("eval_count", 0),
            },
            latency_ms=latency,
            raw_response=data,
        )

    async def generate_async(self, messages: List[Message], **kwargs) -> LLMResponse:
        import time

        import httpx

        start = time.time()

        async with httpx.AsyncClient() as client:
            response = await client.post(
                f"{self.base_url}/api/chat",
                json={
                    "model": self.config.model,
                    "messages": [{"role": m.role, "content": m.content} for m in messages],
                    "stream": False,
                    "options": {
                        "temperature": kwargs.get("temperature", self.config.temperature),
                        "num_predict": kwargs.get("max_tokens", self.config.max_tokens),
                    },
                },
                timeout=self.config.timeout,
            )
            response.raise_for_status()
            data = response.json()

        latency = int((time.time() - start) * 1000)

        return LLMResponse(
            content=data["message"]["content"],
            model=self.config.model,
            provider=LLMProvider.OLLAMA,
            usage={
                "prompt_tokens": data.get("prompt_eval_count", 0),
                "completion_tokens": data.get("eval_count", 0),
                "total_tokens": data.get("prompt_eval_count", 0) + data.get("eval_count", 0),
            },
            latency_ms=latency,
            raw_response=data,
        )


def create_llm(
    provider: Union[str, LLMProvider] = "openai",
    model: Optional[str] = None,
    api_key: Optional[str] = None,
    **kwargs,
) -> BaseLLM:
    """
    Factory function to create an LLM instance.

    Usage:
        llm = create_llm("openai", model="gpt-4o-mini")
        llm = create_llm("anthropic", model="claude-3-sonnet-20240229")
        llm = create_llm("ollama", model="llama3")
    """
    if isinstance(provider, str):
        provider = LLMProvider(provider.lower())

    # Default models per provider
    default_models = {
        LLMProvider.OPENAI: "gpt-4o-mini",
        LLMProvider.ANTHROPIC: "claude-3-sonnet-20240229",
        LLMProvider.OLLAMA: "llama3",
        LLMProvider.GROQ: "llama-3.1-70b-versatile",
        LLMProvider.TOGETHER: "meta-llama/Llama-3-70b-chat-hf",
    }

    config = LLMConfig(
        provider=provider,
        model=model or default_models.get(provider, "gpt-4o-mini"),
        api_key=api_key,
        **kwargs,
    )

    llm_classes = {
        LLMProvider.OPENAI: OpenAILLM,
        LLMProvider.ANTHROPIC: AnthropicLLM,
        LLMProvider.OLLAMA: OllamaLLM,
        # Groq, Together, OpenRouter use OpenAI-compatible API
        LLMProvider.GROQ: OpenAILLM,
        LLMProvider.TOGETHER: OpenAILLM,
        LLMProvider.OPENROUTER: OpenAILLM,
    }

    # Set base_url for OpenAI-compatible providers
    if provider == LLMProvider.GROQ:
        config.base_url = "https://api.groq.com/openai/v1"
        config.api_key = api_key or os.getenv("GROQ_API_KEY")
    elif provider == LLMProvider.TOGETHER:
        config.base_url = "https://api.together.xyz/v1"
        config.api_key = api_key or os.getenv("TOGETHER_API_KEY")
    elif provider == LLMProvider.OPENROUTER:
        config.base_url = "https://openrouter.ai/api/v1"
        config.api_key = api_key or os.getenv("OPENROUTER_API_KEY")

    llm_class = llm_classes.get(provider)
    if llm_class is None:
        raise ValueError(f"Unsupported provider: {provider}")

    return llm_class(config)
