from __future__ import annotations

from typing import Any

import httpx

from ._constants import (
    HEADER_RESP_BUDGET_REMAINING_TOKENS,
    HEADER_RESP_BUDGET_REMAINING_USD,
    HEADER_RESP_COST,
    HEADER_RESP_FAILOVER,
    HEADER_RESP_GOVERNANCE,
    HEADER_RESP_MODEL,
    HEADER_RESP_PROVIDER,
    HEADER_RESP_RUN_ID,
    HEADER_RESP_STEP_ID,
    HEADER_RESP_TOKENS,
    HEADER_RESP_WARNING,
)
from .types.rungate_metadata import RungateMetadata


class RungateResponse:
    """Wraps an httpx response with parsed Rungate metadata."""

    def __init__(self, http_response: httpx.Response) -> None:
        self._response = http_response
        self._body: dict[str, Any] | None = None
        self.metadata = self._parse_metadata()

    @property
    def status_code(self) -> int:
        return self._response.status_code

    @property
    def body(self) -> dict[str, Any]:
        if self._body is None:
            try:
                self._body = self._response.json()
            except Exception:
                # Non-JSON response — return raw text wrapped in a dict
                self._body = {"_raw": self._response.text}
        return self._body

    @property
    def headers(self) -> httpx.Headers:
        return self._response.headers

    def _parse_metadata(self) -> RungateMetadata:
        """Extract Rungate headers into typed metadata."""
        h = self._response.headers
        return RungateMetadata(
            run_id=h.get(HEADER_RESP_RUN_ID),
            step_id=h.get(HEADER_RESP_STEP_ID),
            provider=h.get(HEADER_RESP_PROVIDER),
            model=h.get(HEADER_RESP_MODEL),
            tokens_used=_parse_int(h.get(HEADER_RESP_TOKENS)),
            cost_usd=_parse_float(h.get(HEADER_RESP_COST)),
            budget_remaining_usd=_parse_float(h.get(HEADER_RESP_BUDGET_REMAINING_USD)),
            budget_remaining_tokens=_parse_int(h.get(HEADER_RESP_BUDGET_REMAINING_TOKENS)),
            failover_count=_parse_int(h.get(HEADER_RESP_FAILOVER)),
            warning=h.get(HEADER_RESP_WARNING),
            governance=h.get(HEADER_RESP_GOVERNANCE),
        )


def _parse_int(value: str | None) -> int | None:
    if value is None:
        return None
    try:
        return int(value)
    except (ValueError, TypeError):
        return None


def _parse_float(value: str | None) -> float | None:
    if value is None:
        return None
    try:
        return float(value)
    except (ValueError, TypeError):
        return None
