"""Org plan cache helpers for gateway billing."""

from __future__ import annotations

import logging
import threading
import time
from typing import Any, Callable, Dict, Optional

logger = logging.getLogger(__name__)

FetchFn = Optional[Callable[[str], Optional[Dict[str, Any]]]]


class OrgPlanCache:
    """Simple TTL cache for organization billing configuration."""

    def __init__(
        self,
        fetch_fn: FetchFn,
        *,
        default_tier: str,
        default_price_book: str = "p0-0",
        default_billing_mode: str = "resell",
        ttl_seconds: int = 60,
    ) -> None:
        self._fetch_fn = fetch_fn
        self._default_tier = default_tier
        self._default_price_book = default_price_book
        self._default_billing_mode = default_billing_mode
        self._ttl = ttl_seconds
        self._lock = threading.Lock()
        self._cache: dict[str, tuple[float, Dict[str, Any]]] = {}

    def get(self, org_id: str, *, fallback_tier: Optional[str] = None) -> Dict[str, Any]:
        now = time.monotonic()
        with self._lock:
            cached = self._cache.get(org_id)
            if cached and now - cached[0] < self._ttl:
                return cached[1]

        data: Optional[Dict[str, Any]] = None
        if self._fetch_fn is not None:
            try:
                data = self._fetch_fn(org_id)
            except Exception as exc:
                raise RuntimeError(f"Failed to fetch billing plan for org {org_id!r}.") from exc

        if not data:
            logger.warning(
                "org_plan_fabricated org_id=%s using defaults"
                " (prompt_budget=128000, completion_budget=64000,"
                " billing_mode=%s)",
                org_id,
                self._default_billing_mode,
            )
            data = {}

        plan = {
            "plan_id": data.get("plan_id", "default"),
            "tier": data.get("tier", fallback_tier or self._default_tier),
            "price_book_version": data.get("price_book_version", self._default_price_book),
            "default_billing_mode": data.get("default_billing_mode", self._default_billing_mode),
            "provider_overrides": data.get(
                "provider_overrides", {"takumi": self._default_billing_mode}
            ),
            "prompt_token_budget": int(data.get("prompt_token_budget", 128_000)),
            "completion_token_budget": int(data.get("completion_token_budget", 64_000)),
        }

        overrides = plan.get("provider_overrides")
        if not isinstance(overrides, dict):
            plan["provider_overrides"] = {"takumi": self._default_billing_mode}
        else:
            overrides.setdefault("takumi", self._default_billing_mode)

        with self._lock:
            self._cache[org_id] = (now, plan)
        return plan

    def clear(self) -> None:
        with self._lock:
            self._cache.clear()


__all__ = ["OrgPlanCache"]
