"""Abstract AI provider interface and shared types."""

from __future__ import annotations

import asyncio
import logging
from abc import ABC, abstractmethod
from collections.abc import AsyncGenerator
from dataclasses import dataclass, field
from typing import Any

import httpx

logger = logging.getLogger(__name__)

MAX_RETRIES = 3
RETRY_BASE_DELAY = 1.0  # seconds


async def retry_on_transient(func, *args, **kwargs):
    """Retry an async function on transient HTTP errors (429, 503) with exponential backoff."""
    last_exc = None
    for attempt in range(MAX_RETRIES + 1):
        try:
            return await func(*args, **kwargs)
        except httpx.HTTPStatusError as e:
            last_exc = e
            if e.response.status_code in (429, 503) and attempt < MAX_RETRIES:
                # Check for Retry-After header
                retry_after = e.response.headers.get("retry-after")
                if retry_after:
                    try:
                        delay = min(float(retry_after), 30.0)
                    except ValueError:
                        delay = RETRY_BASE_DELAY * (2 ** attempt)
                else:
                    delay = RETRY_BASE_DELAY * (2 ** attempt)
                logger.warning(
                    "Retrying after %s (attempt %d/%d, status %d)",
                    f"{delay:.1f}s", attempt + 1, MAX_RETRIES, e.response.status_code,
                )
                await asyncio.sleep(delay)
            else:
                raise
    raise last_exc  # type: ignore[misc]


@dataclass(slots=True)
class MediaItem:
    """An image or other media attached to a message."""
    type: str  # "image"
    data: bytes
    mime_type: str  # "image/png", "image/jpeg", etc.


@dataclass(slots=True)
class Message:
    role: str  # "system" | "user" | "assistant" | "tool"
    content: str
    tool_calls: list[ToolCall] = field(default_factory=list)
    tool_call_id: str | None = None
    media: list[MediaItem] = field(default_factory=list)


@dataclass(slots=True)
class ToolDef:
    name: str
    description: str
    parameters: dict[str, Any]


@dataclass(slots=True)
class ToolCall:
    id: str
    name: str
    arguments: dict[str, Any]


@dataclass(slots=True)
class AIResponse:
    content: str
    tool_calls: list[ToolCall]
    usage: dict[str, int]  # {"input": N, "output": N}
    stop_reason: str = ""  # "end_turn", "tool_use", "max_tokens"


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

    All providers normalize to the same Message/AIResponse/ToolCall types.
    """

    name: str

    @abstractmethod
    async def chat(
        self,
        messages: list[Message],
        tools: list[ToolDef] | None = None,
    ) -> AIResponse:
        """Send messages and get a response, optionally with tool use."""
        ...

    async def chat_stream(
        self,
        messages: list[Message],
        tools: list[ToolDef] | None = None,
    ) -> AsyncGenerator[str | AIResponse, None]:
        """Stream text deltas, then yield a final AIResponse.

        Default: falls back to non-streaming chat().
        Providers override for true streaming.
        """
        response = await self.chat(messages, tools)
        if response.content:
            yield response.content
        yield response

    async def close(self) -> None:
        """Close any HTTP clients held by this provider."""
