"""Base model interface for AI models.

This module defines the abstract interface that all AI model implementations
must follow, allowing the orchestrator to work with any AI provider.
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional


class StopReason(str, Enum):
    """Reason the model stopped generating."""

    END_TURN = "end_turn"
    TOOL_USE = "tool_use"
    MAX_TOKENS = "max_tokens"
    STOP_SEQUENCE = "stop_sequence"
    ERROR = "error"


@dataclass
class ToolUse:
    """Represents a tool use request from the model."""

    tool_use_id: str
    name: str
    input: Dict[str, Any]


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

    role: str  # "user", "assistant", "system"
    content: str
    tool_use: Optional[ToolUse] = None
    tool_result: Optional[Dict[str, Any]] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_bedrock_format(self) -> Dict[str, Any]:
        """Convert to Bedrock Converse API format."""
        if self.tool_result:
            return {
                "role": "user",
                "content": [{"toolResult": self.tool_result}],
            }
        if self.tool_use:
            return {
                "role": "assistant",
                "content": [
                    {
                        "toolUse": {
                            "toolUseId": self.tool_use.tool_use_id,
                            "name": self.tool_use.name,
                            "input": self.tool_use.input,
                        }
                    }
                ],
            }
        return {
            "role": self.role,
            "content": [{"text": self.content}],
        }


@dataclass
class ModelResponse:
    """Response from an AI model."""

    content: str
    stop_reason: StopReason
    tool_uses: List[ToolUse] = field(default_factory=list)
    usage: Dict[str, int] = field(default_factory=dict)
    model_id: str = ""
    raw_response: Optional[Dict[str, Any]] = None

    @property
    def has_tool_use(self) -> bool:
        """Check if the response contains tool use requests."""
        return self.stop_reason == StopReason.TOOL_USE and len(self.tool_uses) > 0


@dataclass
class ToolDefinition:
    """Definition of a tool that can be used by the model."""

    name: str
    description: str
    input_schema: Dict[str, Any]

    def to_bedrock_format(self) -> Dict[str, Any]:
        """Convert to Bedrock Converse API format."""
        return {
            "toolSpec": {
                "name": self.name,
                "description": self.description,
                "inputSchema": {"json": self.input_schema},
            }
        }


class BaseModel(ABC):
    """Abstract base class for AI model implementations.

    All AI model implementations (Bedrock, Anthropic, OpenAI, etc.) must
    inherit from this class and implement the required methods.

    Example:
        class MyCustomModel(BaseModel):
            def converse(self, messages, system_prompt, tools):
                # Call your AI API
                return ModelResponse(content="Hello!", stop_reason=StopReason.END_TURN)

            def is_available(self):
                return True
    """

    @abstractmethod
    def converse(
        self,
        messages: List[Message],
        system_prompt: str,
        tools: Optional[List[ToolDefinition]] = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
    ) -> ModelResponse:
        """Send messages to the model and get a response.

        Args:
            messages: Conversation history
            system_prompt: System prompt for the model
            tools: Available tools for the model to use
            max_tokens: Maximum tokens in the response
            temperature: Sampling temperature

        Returns:
            ModelResponse with the model's response

        Raises:
            ModelError: If the model call fails
        """
        pass

    @abstractmethod
    def is_available(self) -> bool:
        """Check if the model is available and configured.

        Returns:
            True if the model can be used, False otherwise
        """
        pass

    @property
    @abstractmethod
    def model_id(self) -> str:
        """Get the model identifier.

        Returns:
            String identifier for the model
        """
        pass

    def supports_tools(self) -> bool:
        """Check if the model supports tool use.

        Returns:
            True if the model supports tools, False otherwise
        """
        return True


class ModelError(Exception):
    """Base exception for model errors."""

    def __init__(
        self,
        message: str,
        model_id: str = "",
        is_transient: bool = False,
        original_error: Optional[Exception] = None,
    ):
        super().__init__(message)
        self.model_id = model_id
        self.is_transient = is_transient
        self.original_error = original_error


class ModelNotAvailableError(ModelError):
    """Raised when a model is not available."""

    pass


class ModelRateLimitError(ModelError):
    """Raised when rate limited by the model provider."""

    def __init__(self, message: str, retry_after: Optional[int] = None, **kwargs):
        super().__init__(message, is_transient=True, **kwargs)
        self.retry_after = retry_after
