"""
LLM Interface Module

Abstract interface for language model providers:
- Model-agnostic completion API
- Streaming support
- Token counting
- Error handling with retries
"""

import json
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, AsyncIterator, Callable, Optional


@dataclass
class Message:
    """A message in the conversation."""

    role: str  # "system", "user", "assistant", "tool"
    content: str
    tool_call_id: Optional[str] = None
    tool_calls: Optional[list[dict]] = None
    name: Optional[str] = None  # For tool messages


@dataclass
class ToolCall:
    """A tool call from the LLM."""

    id: str
    name: str
    arguments: dict[str, Any]

    @classmethod
    def from_dict(cls, data: dict) -> "ToolCall":
        """Create from API response dict."""
        args = data.get("arguments", {})
        if isinstance(args, str):
            args = json.loads(args)
        return cls(
            id=data.get("id", ""),
            name=data.get("name", data.get("function", {}).get("name", "")),
            arguments=args,
        )


@dataclass
class CompletionResponse:
    """Response from an LLM completion."""

    content: Optional[str] = None
    tool_calls: list[ToolCall] = field(default_factory=list)
    finish_reason: str = "stop"
    usage: Optional[dict[str, int]] = None
    raw_response: Optional[Any] = None

    @property
    def has_tool_calls(self) -> bool:
        """Check if response contains tool calls."""
        return len(self.tool_calls) > 0


@dataclass
class LLMConfig:
    """Configuration for LLM provider."""

    model: str = "gpt-4o"
    max_tokens: int = 4096
    temperature: float = 0.1
    api_key: Optional[str] = None
    api_base: Optional[str] = None
    timeout: int = 120
    max_retries: int = 3


class LLMProvider(ABC):
    """
    Abstract base class for LLM providers.
    
    Implementations must provide:
    - complete(): Synchronous completion
    - complete_stream(): Streaming completion (optional)
    """

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

    @abstractmethod
    def complete(
        self,
        messages: list[Message],
        tools: Optional[list[dict]] = None,
    ) -> CompletionResponse:
        """
        Generate a completion from the messages.
        
        Args:
            messages: List of conversation messages
            tools: Optional list of tool definitions
            
        Returns:
            CompletionResponse with content and/or tool calls
        """
        pass

    def complete_stream(
        self,
        messages: list[Message],
        tools: Optional[list[dict]] = None,
    ) -> AsyncIterator[str]:
        """
        Generate a streaming completion.
        
        Default implementation falls back to non-streaming.
        """
        raise NotImplementedError("Streaming not supported by this provider")

    def count_tokens(self, text: str) -> int:
        """
        Estimate token count for text.
        
        Default implementation uses approximate character ratio.
        """
        # Rough approximation: ~4 chars per token for English
        return len(text) // 4

    def format_messages(self, messages: list[Message]) -> list[dict]:
        """Convert Message objects to API format."""
        return [self._format_message(m) for m in messages]

    def _format_message(self, message: Message) -> dict:
        """Format a single message for the API."""
        msg = {"role": message.role, "content": message.content}

        if message.tool_call_id:
            msg["tool_call_id"] = message.tool_call_id

        if message.tool_calls:
            msg["tool_calls"] = message.tool_calls

        if message.name:
            msg["name"] = message.name

        return msg


class LLMError(Exception):
    """Base exception for LLM errors."""

    def __init__(self, message: str, retriable: bool = False):
        super().__init__(message)
        self.retriable = retriable


class RateLimitError(LLMError):
    """Rate limit exceeded error."""

    def __init__(self, message: str, retry_after: Optional[float] = None):
        super().__init__(message, retriable=True)
        self.retry_after = retry_after


class AuthenticationError(LLMError):
    """API authentication failed."""

    def __init__(self, message: str):
        super().__init__(message, retriable=False)


class ModelNotFoundError(LLMError):
    """Requested model not available."""

    def __init__(self, message: str):
        super().__init__(message, retriable=False)


# Provider registry
_providers: dict[str, type[LLMProvider]] = {}


def register_provider(name: str) -> Callable[[type[LLMProvider]], type[LLMProvider]]:
    """Decorator to register an LLM provider."""

    def decorator(cls: type[LLMProvider]) -> type[LLMProvider]:
        _providers[name] = cls
        return cls

    return decorator


def get_provider(name: str, config: LLMConfig) -> LLMProvider:
    """Get an LLM provider instance by name."""
    if name not in _providers:
        available = ", ".join(_providers.keys())
        raise ValueError(f"Unknown provider: {name}. Available: {available}")

    return _providers[name](config)


def list_providers() -> list[str]:
    """List available provider names."""
    return list(_providers.keys())
