"""Abstract provider bridge interface."""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import AsyncIterator

import aiohttp

from backend.config import ModelConfig, ProviderConfig


@dataclass
class BridgeRequest:
    """Normalized request to send to a provider."""
    messages: list[dict]
    model: str
    temperature: float | None = None
    max_tokens: int | None = None
    top_p: float | None = None
    stream: bool = False
    tools: list[dict] | None = None
    tool_choice: str | dict | None = None
    stop: list[str] | str | None = None
    # Original request body for passthrough fields
    extra: dict = field(default_factory=dict)


@dataclass
class BridgeResponse:
    """Normalized response from a provider."""
    content: str
    model: str
    finish_reason: str | None = None
    tokens_in: int = 0
    tokens_out: int = 0
    # Anthropic cache tokens (prompt caching)
    cache_read_tokens: int = 0
    cache_creation_tokens: int = 0
    tool_calls: list[dict] | None = None
    raw: dict | None = None


@dataclass
class StreamUsage:
    """Usage stats accumulated during streaming."""
    tokens_in: int = 0
    tokens_out: int = 0
    cache_read_tokens: int = 0
    cache_creation_tokens: int = 0


class StreamResult:
    """Wrapper for streaming response that captures usage.
    
    Usage is populated during iteration and available after stream completes.
    
    Example:
        result = await bridge.stream(request, model)
        async for line in result:
            # process line
        print(result.usage.tokens_out)
    """
    
    def __init__(self, iterator: AsyncIterator[str]):
        self._iterator = iterator
        self.usage = StreamUsage()
    
    def __aiter__(self):
        return self
    
    async def __anext__(self) -> str:
        return await self._iterator.__anext__()


class BaseBridge(ABC):
    """Abstract bridge for translating between OpenAI format and provider-specific APIs."""

    def __init__(self, provider: ProviderConfig, session: aiohttp.ClientSession):
        self.provider = provider
        self.session = session

    @abstractmethod
    async def complete(self, request: BridgeRequest, model: ModelConfig) -> BridgeResponse:
        """Send a non-streaming completion request."""
        ...

    @abstractmethod
    async def stream(self, request: BridgeRequest, model: ModelConfig) -> StreamResult:
        """Send a streaming completion request.
        
        Returns a StreamResult that yields SSE lines and accumulates usage stats.
        """
        ...

    @property
    @abstractmethod
    def provider_type(self) -> str:
        """Provider type identifier (openai, anthropic, bedrock)."""
        ...
