"""
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).

    ``/chat`` returns ``{content: "..."}``; ``/generate`` returns ``{text: "..."}``.
    Both are accessible via ``.content``.
    """

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

    @property
    def content(self) -> str:
        """The generated text (works for both /chat and /generate responses)."""
        # /chat → "content"; /generate → "text"
        return self._data.get("content") or self._data.get("text") or ""

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


class StreamChunk(_APIObject):
    """A single SSE chunk from the native streaming API.

    The native API sends ``{"token": "...", "model": "..."}`` per chunk.
    Final chunk may include usage stats: ``{"done": true, "prompt_tokens": ..., ...}``.
    """

    done: bool
    token: Optional[str]
    model: Optional[str]
    prompt_tokens: Optional[int]
    completion_tokens: Optional[int]
    tokens_per_second: Optional[float]

    def __str__(self) -> str:
        """Return the token text for this chunk."""
        return self._data.get("token") or ""


class StreamResponse:
    """Iterable wrapper over a native SSE stream.

    Iterate to receive ``StreamChunk`` objects in real-time. After iteration
    completes, access ``.content`` for the full concatenated output.

    Example::

        stream = client.chat.create(model="qwen3:8b",
                                    messages=[...], stream=True)
        for chunk in stream:
            print(str(chunk), end="", flush=True)
        print()
        print("Full:", stream.content)
    """

    def __init__(self, chunks_iter: Iterator[StreamChunk]):
        self._iter = chunks_iter
        self._collected: List[str] = []
        self._prompt_tokens: int = 0
        self._completion_tokens: int = 0
        self._tokens_per_second: float = 0.0

    def __iter__(self) -> Iterator[StreamChunk]:
        for chunk in self._iter:
            token = chunk._data.get("token")
            if token:
                self._collected.append(token)
            # Collect usage from any chunk (last one usually has it)
            if chunk._data.get("prompt_tokens"):
                self._prompt_tokens = chunk._data["prompt_tokens"]
            if chunk._data.get("completion_tokens"):
                self._completion_tokens = chunk._data["completion_tokens"]
            if chunk._data.get("tokens_per_second"):
                self._tokens_per_second = chunk._data["tokens_per_second"]
            yield chunk

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

    @property
    def prompt_tokens(self) -> int:
        return self._prompt_tokens

    @property
    def completion_tokens(self) -> int:
        return self._completion_tokens

    @property
    def tokens_per_second(self) -> float:
        return self._tokens_per_second

    @property
    def usage(self) -> Dict[str, Any]:
        return {
            "prompt_tokens": self._prompt_tokens,
            "completion_tokens": self._completion_tokens,
            "total_tokens": self._prompt_tokens + self._completion_tokens,
            "tokens_per_second": self._tokens_per_second,
        }


# StreamDelta kept for backward compatibility
class StreamDelta(_APIObject):
    """Backward-compatible stub (native API uses StreamChunk.token directly)."""

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

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


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

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

    name: str
    display_name: str
    context_length: int
    model_type: str
    is_active: bool
    input_price_per_million: float
    output_price_per_million: float

    @property
    def id(self) -> str:
        """Alias for ``name`` (model identifier)."""
        return self._data.get("name") or self._data.get("id") or ""


class ModelList(_APIObject):
    """List of models (native API format: ``{models: [...]}``).

    Also handles legacy ``{data: [...]}`` format.
    """

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

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

    # Backward-compat alias
    @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."""
