"""LLM client abstraction for multiple providers."""

import os
from enum import Enum
from typing import Optional, List, Dict, Any
from dataclasses import dataclass


class LLMProvider(Enum):
    """Supported LLM providers."""
    OPENAI = "openai"
    ANTHROPIC = "anthropic"


@dataclass
class LLMMessage:
    """A message in LLM conversation."""
    role: str  # "system", "user", or "assistant"
    content: str


class LLMClient:
    """Unified client for LLM providers.

    Supports OpenAI and Anthropic APIs with BYOK (Bring Your Own Key).
    """

    def __init__(self, provider: Optional[LLMProvider] = None, api_key: Optional[str] = None):
        """Initialize LLM client.

        Args:
            provider: LLM provider to use (auto-detected if None)
            api_key: API key (auto-detected from env if None)
        """
        self.provider = provider or self._detect_provider()
        self.api_key = api_key or self._get_api_key()

        if not self.api_key:
            raise ValueError(
                f"No API key found for {self.provider.value}. "
                f"Set {self._get_env_var_name()} environment variable."
            )

        self._client = self._initialize_client()

    def _detect_provider(self) -> LLMProvider:
        """Auto-detect which provider to use based on environment."""
        if os.getenv("ANTHROPIC_API_KEY"):
            return LLMProvider.ANTHROPIC
        elif os.getenv("OPENAI_API_KEY"):
            return LLMProvider.OPENAI
        else:
            raise ValueError(
                "No LLM API key found. Set OPENAI_API_KEY or ANTHROPIC_API_KEY environment variable."
            )

    def _get_env_var_name(self) -> str:
        """Get the environment variable name for the current provider."""
        if self.provider == LLMProvider.OPENAI:
            return "OPENAI_API_KEY"
        elif self.provider == LLMProvider.ANTHROPIC:
            return "ANTHROPIC_API_KEY"
        else:
            raise ValueError(f"Unknown provider: {self.provider}")

    def _get_api_key(self) -> Optional[str]:
        """Get API key from environment."""
        return os.getenv(self._get_env_var_name())

    def _initialize_client(self) -> Any:
        """Initialize the provider-specific client."""
        if self.provider == LLMProvider.OPENAI:
            from openai import OpenAI
            return OpenAI(api_key=self.api_key)
        elif self.provider == LLMProvider.ANTHROPIC:
            from anthropic import Anthropic
            return Anthropic(api_key=self.api_key)
        else:
            raise ValueError(f"Unsupported provider: {self.provider}")

    def complete(
        self,
        messages: List[LLMMessage],
        model: Optional[str] = None,
        temperature: float = 0.7,
        max_tokens: int = 2000
    ) -> str:
        """Get a completion from the LLM.

        Args:
            messages: Conversation messages
            model: Model name (uses default if None)
            temperature: Sampling temperature (0.0-1.0)
            max_tokens: Maximum tokens in response

        Returns:
            The LLM's response text
        """
        if self.provider == LLMProvider.OPENAI:
            return self._complete_openai(messages, model, temperature, max_tokens)
        elif self.provider == LLMProvider.ANTHROPIC:
            return self._complete_anthropic(messages, model, temperature, max_tokens)
        else:
            raise ValueError(f"Unsupported provider: {self.provider}")

    def _complete_openai(
        self,
        messages: List[LLMMessage],
        model: Optional[str],
        temperature: float,
        max_tokens: int
    ) -> str:
        """Complete using OpenAI API."""
        model = model or "gpt-4-turbo-preview"

        response = self._client.chat.completions.create(
            model=model,
            messages=[{"role": m.role, "content": m.content} for m in messages],
            temperature=temperature,
            max_tokens=max_tokens
        )

        return response.choices[0].message.content

    def _complete_anthropic(
        self,
        messages: List[LLMMessage],
        model: Optional[str],
        temperature: float,
        max_tokens: int
    ) -> str:
        """Complete using Anthropic API."""
        model = model or "claude-3-5-sonnet-20241022"

        # Anthropic API requires system message separate from messages
        system_messages = [m.content for m in messages if m.role == "system"]
        system = "\n\n".join(system_messages) if system_messages else None

        conversation_messages = [
            {"role": m.role, "content": m.content}
            for m in messages
            if m.role != "system"
        ]

        response = self._client.messages.create(
            model=model,
            max_tokens=max_tokens,
            temperature=temperature,
            system=system,
            messages=conversation_messages
        )

        return response.content[0].text
