"""
AgniPod SDK — Response Types
==============================

Lightweight dataclass-style wrappers around raw API JSON.
Every type supports attribute AND dict-style access for convenience.
"""

from __future__ import annotations
from typing import Any, Dict, Iterator, List, Optional


class _APIObject:
    """Base mixin: stores raw JSON and exposes attribute access."""

    _data: Dict[str, Any]

    def __init__(self, data: Dict[str, Any] | None = None, **kwargs):
        object.__setattr__(self, "_data", {**(data or {}), **kwargs})

    def __getattr__(self, name: str) -> Any:
        try:
            return self._data[name]
        except KeyError:
            raise AttributeError(
                f"{type(self).__name__!r} has no attribute {name!r}"
            )

    def __getitem__(self, key: str) -> Any:
        return self._data[key]

    def __contains__(self, key: str) -> bool:
        return key in self._data

    def get(self, key: str, default: Any = None) -> Any:
        return self._data.get(key, default)

    def to_dict(self) -> Dict[str, Any]:
        """Return the raw API response as a plain dict."""
        return dict(self._data)

    def __repr__(self) -> str:
        fields = ", ".join(
            f"{k}={v!r}"
            for k, v in self._data.items()
            if not k.startswith("_") and v is not None
        )
        return f"{type(self).__name__}({fields})"

    def __bool__(self) -> bool:
        return bool(self._data)


# ── Inference ────────────────────────────────────────────────────────────

class Completion(_APIObject):
    """Response from ``/generate`` or ``/chat`` (non-streaming)."""

    id: str
    model: str
    content: str
    prompt_tokens: int
    completion_tokens: int
    tokens_per_second: float

    def __str__(self) -> str:
        return self.content


class StreamDelta(_APIObject):
    """A single delta within a streaming chunk."""

    role: Optional[str]
    content: Optional[str]

    def __str__(self) -> str:
        return self.get("content") or ""


class StreamChunk(_APIObject):
    """A single SSE chunk during streaming."""

    id: str
    done: bool
    delta: Optional[StreamDelta]
    usage: Optional[Dict[str, int]]

    def __init__(self, data: Dict[str, Any] | None = None, **kwargs):
        super().__init__(data, **kwargs)
        raw_delta = self._data.get("delta")
        if isinstance(raw_delta, dict):
            object.__setattr__(self, "_delta_obj", StreamDelta(raw_delta))
        else:
            object.__setattr__(self, "_delta_obj", None)

    @property
    def delta(self) -> Optional[StreamDelta]:
        return self._delta_obj

    def __str__(self) -> str:
        return str(self._delta_obj) if self._delta_obj else ""


class StreamResponse:
    """Iterable wrapper over an SSE stream.

    Iterate to receive chunks in real-time, then access ``.content``
    for the fully concatenated output once iteration completes.
    """

    def __init__(self, chunks_iter: Iterator[StreamChunk]):
        self._iter = chunks_iter
        self._collected_content: List[str] = []
        self._final_usage: Optional[Dict[str, int]] = None
        self._final_id: Optional[str] = None

    def __iter__(self) -> Iterator[StreamChunk]:
        for chunk in self._iter:
            if chunk.delta and chunk.delta.get("content"):
                self._collected_content.append(chunk.delta.content)
            if chunk.get("id"):
                self._final_id = chunk.id
            if chunk.get("usage"):
                self._final_usage = chunk.usage
            yield chunk

    @property
    def content(self) -> str:
        """Full concatenated content (available after iteration)."""
        return "".join(self._collected_content)

    @property
    def usage(self) -> Optional[Dict[str, int]]:
        return self._final_usage


# ── Models ───────────────────────────────────────────────────────────────

class Model(_APIObject):
    """A single model record."""

    id: str
    object: str
    created: int
    owned_by: str


class ModelList(_APIObject):
    """List of models."""

    def __init__(self, data: Dict[str, Any] | None = None, **kwargs):
        super().__init__(data, **kwargs)
        self._models = [Model(m) for m in self._data.get("data", [])]

    @property
    def data(self) -> List[Model]:
        return self._models

    def __iter__(self) -> Iterator[Model]:
        return iter(self._models)

    def __len__(self) -> int:
        return len(self._models)


# ── Batches ──────────────────────────────────────────────────────────────

class Batch(_APIObject):
    """A batch processing job."""

    id: int
    status: str
    total_items: int
    processed_items: int
    failed_items: int


class BatchCreateResult(_APIObject):
    """Result from creating a batch."""

    batch_id: int
    total_items: int
    is_valid: bool


class BatchProgress(_APIObject):
    """Progress of a running batch."""

    id: int
    status: str
    total_items: int
    processed_items: int
    failed_items: int
    progress_percent: float


class BatchList(_APIObject):
    """Paginated list of batches."""

    def __init__(self, data: Dict[str, Any] | None = None, **kwargs):
        super().__init__(data, **kwargs)
        self._batches = [Batch(b) for b in self._data.get("batches", [])]

    @property
    def batches(self) -> List[Batch]:
        return self._batches

    def __iter__(self) -> Iterator[Batch]:
        return iter(self._batches)

    def __len__(self) -> int:
        return len(self._batches)


# ── API Keys ─────────────────────────────────────────────────────────────

class ApiKey(_APIObject):
    """An API key record."""

    id: int
    key_prefix: str
    name: Optional[str]
    is_active: bool


class ApiKeyCreated(_APIObject):
    """Newly created API key (includes the full secret)."""

    id: int
    key: str
    key_prefix: str


class ApiKeyList(_APIObject):
    """Paginated list of API keys."""

    def __init__(self, data: Dict[str, Any] | None = None, **kwargs):
        super().__init__(data, **kwargs)
        self._keys = [ApiKey(k) for k in self._data.get("keys", [])]

    @property
    def keys(self) -> List[ApiKey]:
        return self._keys

    def __iter__(self) -> Iterator[ApiKey]:
        return iter(self._keys)

    def __len__(self) -> int:
        return len(self._keys)


class KeyUsage(_APIObject):
    """Usage statistics for an API key."""

    key_id: int
    total_requests: int
    total_tokens: int
    total_cost: float


# ── Billing ──────────────────────────────────────────────────────────────

class Balance(_APIObject):
    """Wallet balance."""

    balance: float
    currency: str

    def __str__(self) -> str:
        return f"${self.balance:.4f} {self.currency}"


class Wallet(_APIObject):
    """Full wallet details."""


class Transaction(_APIObject):
    """A single transaction record."""

    id: int
    transaction_type: str
    amount: float
    balance_after: float


class TransactionList(_APIObject):
    """Paginated list of transactions."""

    def __init__(self, data: Dict[str, Any] | None = None, **kwargs):
        super().__init__(data, **kwargs)
        self._txns = [Transaction(t) for t in self._data.get("transactions", [])]

    @property
    def transactions(self) -> List[Transaction]:
        return self._txns

    def __iter__(self) -> Iterator[Transaction]:
        return iter(self._txns)

    def __len__(self) -> int:
        return len(self._txns)


class UsageSummary(_APIObject):
    """Aggregated usage over a period."""
