import asyncio
import random
import re
from abc import ABC, abstractmethod
from collections.abc import AsyncIterator
from dataclasses import dataclass, field
from typing import Any

from otto.errors import OttoExecutionError

# === Errors ===


class OAuthProviderError(OttoExecutionError):
    """Base exception for OAuth provider errors."""


@dataclass
class _Function:
    name: str | None = None
    arguments: str | None = None


@dataclass
class _ToolCallDelta:
    index: int = 0
    id: str | None = None
    function: _Function | None = None
    thought_signature: str | None = None


@dataclass
class _Delta:
    content: str | None = None
    role: str | None = None
    tool_calls: list[_ToolCallDelta] | None = None


@dataclass
class _Choice:
    delta: _Delta = field(default_factory=_Delta)
    finish_reason: str | None = None
    index: int = 0


@dataclass
class _Usage:
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0


@dataclass
class ProviderChunk:
    choices: list[_Choice] = field(default_factory=lambda: [_Choice()])
    usage: _Usage | None = None


# === Base Provider ===


class OAuthProvider(ABC):
    """Abstract base class for OAuth-authenticated providers."""

    provider_name: str = "base"
    model_prefixes: list[str] = []
    system_prompt_prefix: str | None = None

    def __init__(self, credentials: dict[str, Any]):
        self._credentials = credentials

    @abstractmethod
    async def stream_completion(
        self,
        model: str,
        messages: list[dict[str, Any]],
        *,
        tools: list[dict[str, Any]] | None = None,
        tool_choice: str | dict[str, Any] | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
        system_prompt: str | None = None,
        session_id: str | None = None,
        reasoning_effort: str | None = None,
    ) -> AsyncIterator[Any]:
        raise NotImplementedError
        yield

    def _text_chunk(self, content: str) -> ProviderChunk:
        return ProviderChunk(choices=[_Choice(delta=_Delta(content=content))])

    def _tool_call_chunk(
        self,
        index: int,
        id: str | None,
        name: str | None,
        arguments: str | None,
        thought_signature: str | None = None,
    ) -> ProviderChunk:
        function = _Function(name=name, arguments=arguments)
        tool_call = _ToolCallDelta(
            index=index,
            id=id,
            function=function,
            thought_signature=thought_signature,
        )
        return ProviderChunk(choices=[_Choice(delta=_Delta(tool_calls=[tool_call]))])

    def _final_chunk(self, finish_reason: str, usage: dict[str, Any] | None) -> ProviderChunk:
        chunk_usage = None
        if usage is not None:
            prompt_tokens = int(usage.get("prompt_tokens", 0) or 0)
            completion_tokens = int(usage.get("completion_tokens", 0) or 0)
            total_tokens = int(usage.get("total_tokens", prompt_tokens + completion_tokens) or 0)
            chunk_usage = _Usage(
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                total_tokens=total_tokens,
            )
        return ProviderChunk(choices=[_Choice(finish_reason=finish_reason)], usage=chunk_usage)

    def supports_model(self, model: str) -> bool:
        """Return True if this provider can handle the model string."""
        model_lower = model.lower()
        return any(
            model_lower == prefix.lower() or model_lower.startswith(prefix.lower() + "/")
            for prefix in self.model_prefixes
        )

    async def _request_with_retry(
        self,
        session: Any,
        method: str,
        url: str,
        headers: dict[str, str],
        json_body: dict[str, Any],
        max_retries: int = 3,
        max_retry_delay: float = 30.0,
    ) -> Any:
        delay = 1.0
        for attempt in range(max_retries + 1):
            response = await session.request(method, url, headers=headers, json=json_body)
            if response.status < 400:
                return response
            if response.status == 429 or response.status >= 500:
                if attempt < max_retries:
                    delay = self._parse_retry_delay(response, delay, max_retry_delay)
                    body_text = await response.text()
                    response.release()
                    body_delay = self._parse_body_retry_hint(body_text)
                    if body_delay is not None:
                        delay = min(body_delay + 0.5, max_retry_delay)
                    jitter = delay * 0.25 * random.random()
                    await asyncio.sleep(delay + jitter)
                    delay = min(delay * 2, max_retry_delay)
                    continue
            body = await response.text()
            response.release()
            raise OAuthProviderError(f"Request failed ({response.status}): {body}")
        raise OAuthProviderError("Request failed after retries")

    @staticmethod
    def _parse_retry_delay(response: Any, current_delay: float, max_delay: float) -> float:
        retry_after = getattr(response, "headers", {}).get("Retry-After")
        if retry_after:
            try:
                return max(min(float(retry_after), max_delay), 0.5)
            except ValueError:
                pass
        return current_delay

    @staticmethod
    def _parse_body_retry_hint(body: str) -> float | None:
        match = re.search(r"reset after (\d+)s", body)
        if match:
            return float(match.group(1))
        match = re.search(r"retry after (\d+)", body, re.IGNORECASE)
        if match:
            return float(match.group(1))
        return None
