"""Provider base class and interfaces for Oclawma."""

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import AsyncIterator
from dataclasses import dataclass
from enum import Enum
from typing import Any, Literal


class ProviderError(Exception):
    """Base exception for provider errors."""

    def __init__(self, message: str, cause: Exception | None = None) -> None:
        super().__init__(message)
        self.message = message
        self.cause = cause


class ProviderConnectionError(ProviderError):
    """Raised when unable to connect to the provider."""


class AuthenticationError(ProviderError):
    """Raised when authentication fails."""


class RateLimitError(ProviderError):
    """Raised when rate limit is exceeded."""


class ModelNotFoundError(ProviderError):
    """Raised when the requested model is not available."""


class CompletionError(ProviderError):
    """Raised when completion generation fails."""


@dataclass(frozen=True)
class Message:
    """A chat message."""

    role: Literal["system", "user", "assistant"]
    content: str
    name: str | None = None


@dataclass(frozen=True)
class CompletionRequest:
    """Request for chat completion."""

    messages: list[Message]
    model: str
    temperature: float = 0.7
    max_tokens: int | None = None
    top_p: float = 1.0
    stream: bool = False


@dataclass(frozen=True)
class CompletionResponse:
    """Response from chat completion."""

    content: str
    model: str
    usage: UsageStats
    finish_reason: str | None = None


@dataclass(frozen=True)
class UsageStats:
    """Token usage statistics."""

    prompt_tokens: int
    completion_tokens: int
    total_tokens: int


class FinishReason(Enum):
    """Reason for completion finishing."""

    STOP = "stop"
    LENGTH = "length"
    CONTENT_FILTER = "content_filter"
    TOOL_CALLS = "tool_calls"


class BaseProvider(ABC):
    """Abstract base class for AI providers."""

    def __init__(self, base_url: str, api_key: str | None = None) -> None:
        """Initialize the provider.

        Args:
            base_url: The base URL for the provider API.
            api_key: Optional API key for authentication.
        """
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key

    @abstractmethod
    async def complete(self, request: CompletionRequest) -> CompletionResponse:
        """Generate a chat completion.

        Args:
            request: The completion request.

        Returns:
            The completion response.

        Raises:
            ConnectionError: If unable to connect to the provider.
            AuthenticationError: If authentication fails.
            ModelNotFoundError: If the model is not available.
            CompletionError: If completion generation fails.
        """
        ...

    @abstractmethod
    def stream_complete(self, request: CompletionRequest) -> AsyncIterator[CompletionResponse]:
        """Generate a streaming chat completion.

        Args:
            request: The completion request. Must have stream=True.

        Yields:
            Completion response chunks.

        Raises:
            ConnectionError: If unable to connect to the provider.
            AuthenticationError: If authentication fails.
            ModelNotFoundError: If the model is not available.
            CompletionError: If completion generation fails.
        """
        ...

    @abstractmethod
    async def list_models(self) -> list[str]:
        """List available models.

        Returns:
            List of available model names.

        Raises:
            ConnectionError: If unable to connect to the provider.
        """
        ...

    @abstractmethod
    def count_tokens(self, text: str, model: str | None = None) -> int:
        """Estimate the number of tokens in the given text.

        For local models without access to the tokenizer, this uses
        a character-based heuristic (approximately 4 characters per token).

        Args:
            text: The text to count tokens for.
            model: Optional model name for more accurate counting.

        Returns:
            Estimated token count.
        """
        ...

    def count_message_tokens(self, messages: list[Message], model: str | None = None) -> int:
        """Estimate tokens for a list of messages.

        Args:
            messages: The messages to count.
            model: Optional model name for more accurate counting.

        Returns:
            Estimated total token count.
        """
        total = 0
        for msg in messages:
            # Base tokens per message (role delimiters, etc.)
            total += 4
            total += self.count_tokens(msg.content, model)
            if msg.name:
                total += self.count_tokens(msg.name, model)
        # Add completion primer tokens
        total += 2
        return total

    @abstractmethod
    async def health_check(self) -> dict[str, Any]:
        """Check the health of the provider connection.

        Returns:
            Health status information.
        """
        ...
