"""Provider-agnostic interface for structured (JSON-schema) model calls (design §7.5, AI-01)."""
from __future__ import annotations
import base64
from dataclasses import dataclass, field
from typing import Protocol, Type

from pydantic import BaseModel


@dataclass
class TextPart:
    text: str


@dataclass
class ImagePart:
    data: bytes
    media_type: str = "image/jpeg"

    def b64(self) -> str:
        return base64.b64encode(self.data).decode()


@dataclass
class Result:
    parsed: BaseModel
    raw_text: str
    provider: str
    model: str
    tokens_in: int = 0
    tokens_out: int = 0
    cached_tokens: int = 0
    latency_ms: int = 0
    cost_micro: int = 0
    extra: dict = field(default_factory=dict)

    def model_info(self) -> dict:
        return {"provider": self.provider, "model": self.model}

    def tokens(self) -> dict:
        return {"in": self.tokens_in, "out": self.tokens_out, "cached": self.cached_tokens}


class ProviderError(Exception):
    """Wraps vendor errors. retryable=True for rate limits, 5xx, timeouts."""

    def __init__(self, msg: str, retryable: bool):
        super().__init__(msg)
        self.retryable = retryable


class SchemaError(Exception):
    """Model output did not validate against the schema. Carries the raw text for the retry prompt."""

    def __init__(self, msg: str, raw_text: str):
        super().__init__(msg)
        self.raw_text = raw_text


class Provider(Protocol):
    name: str

    def structured(self, *, model: str, system: str, parts: list[TextPart | ImagePart],
                   schema: Type[BaseModel], max_tokens: int = 8000,
                   cache_system: bool = True) -> Result: ...


def cost_micro(pricing: dict, model: str, tokens_in: int, tokens_out: int, cached: int) -> int:
    """pricing[model] = {"in": usd per 1M input tokens, "out": ..., "cached_in": ...}."""
    p = pricing.get(model)
    if not p:
        return 0
    usd = ((tokens_in - cached) * p.get("in", 0) + cached * p.get("cached_in", p.get("in", 0))
           + tokens_out * p.get("out", 0)) / 1_000_000
    return int(round(usd * 1_000_000))


def schema_name(schema: Type[BaseModel]) -> str:
    return "submit_" + schema.__name__.lower()