"""Provider Capacity Manager: observe, reserve, wait, release, reconcile."""

from __future__ import annotations

import asyncio
import hashlib
import time
from collections.abc import Mapping
from typing import Any

import structlog

from src.core.provider_capacity.store import (
    CapacityStore,
    MemoryCapacityStore,
    RedisCapacityStore,
)
from src.core.provider_capacity.types import (
    CapacityBucket,
    CapacityObservation,
    CapacityPolicy,
    CapacityReservation,
    ProviderCapacityUnavailableError,
    ReserveDecision,
)

logger = structlog.get_logger()

_memory_fallback = MemoryCapacityStore()
_test_store: CapacityStore | None = None


def override_capacity_store(store: CapacityStore | None) -> None:
    global _test_store
    _test_store = store


def reset_capacity_state() -> None:
    global _memory_fallback
    _memory_fallback = MemoryCapacityStore()
    override_capacity_store(None)


def account_scope_id(secret: str | None) -> str:
    if not secret:
        return "unscoped"
    return hashlib.sha256(secret.encode("utf-8")).hexdigest()[:16]


def policy_from_settings(settings: Any) -> CapacityPolicy:
    return CapacityPolicy(
        safety_fraction=float(getattr(settings, "provider_capacity_safety_fraction", 0.90)),
        wait_seconds=float(getattr(settings, "provider_capacity_wait_seconds", 60)),
        unknown_max_in_flight=int(getattr(settings, "provider_capacity_unknown_max_in_flight", 2)),
        stale_observation_seconds=float(
            getattr(settings, "provider_capacity_stale_observation_seconds", 120)
        ),
    )


def store_from_settings(settings: Any) -> CapacityStore:
    if _test_store is not None:
        return _test_store
    redis_url = getattr(settings, "redis_url", None)
    if redis_url:
        return RedisCapacityStore(redis_url)
    return _memory_fallback


def watchdog_seconds(settings: Any, provider_timeout: float) -> float:
    """Outer ceiling for capacity-wait plus provider execution. Not a provider timeout."""
    if not getattr(settings, "provider_capacity_enabled", True):
        return provider_timeout
    return provider_timeout + float(getattr(settings, "provider_capacity_wait_seconds", 60))


def capacity_bucket(
    *,
    provider: str,
    secret: str | None,
    model: str,
    dimension: str,
) -> CapacityBucket:
    return CapacityBucket(
        provider=provider,
        account_scope=account_scope_id(secret),
        model=model or "unspecified",
        dimension=dimension,
    )


class ProviderCapacityManager:
    def __init__(self, store: CapacityStore, policy: CapacityPolicy):
        self.store = store
        self.policy = policy

    async def acquire(
        self,
        bucket: CapacityBucket,
        demand: int,
        *,
        trace: dict | None = None,
    ) -> CapacityReservation:
        waiter_id = hashlib.sha1(f"{bucket.id}:{time.time_ns()}".encode()).hexdigest()[:20]
        started = time.monotonic()
        queued = False
        last: ReserveDecision | None = None
        used_fallback = False
        demand = max(1, int(demand))
        try:
            while True:
                now_ms = int(time.time() * 1000)
                active_store = self.store
                try:
                    last, snapshot = await self.store.try_reserve(
                        bucket, demand, waiter_id, now_ms, self.policy
                    )
                except Exception:
                    used_fallback = True
                    active_store = _memory_fallback
                    last, snapshot = await active_store.try_reserve(
                        bucket, demand, waiter_id, now_ms, self.policy
                    )
                    logger.warning(
                        "provider_capacity_redis_fallback",
                        provider_capacity_bucket=bucket.id,
                        capacity_dimension=bucket.dimension,
                    )
                record = _trace_fields(bucket, demand, last, snapshot, started)
                if last.granted:
                    record["capacity_granted"] = True
                    record["capacity_queued"] = queued
                    record["capacity_immediate_grant"] = not queued
                    record["capacity_wait_ms"] = int((time.monotonic() - started) * 1000)
                    _bind_trace(trace, record)
                    logger.info("provider_capacity_granted", **record)
                    return CapacityReservation(
                        bucket=bucket,
                        waiter_id=waiter_id,
                        demand=demand,
                        granted_at_monotonic=time.monotonic(),
                        wait_ms=record["capacity_wait_ms"],
                        decision=last,
                        used_fallback=used_fallback,
                    )
                queued = True
                already_logged = bool(
                    (trace or {}).get("provider_capacity", {}).get("capacity_queued")
                )
                if not already_logged:
                    logger.info("provider_capacity_queued", **record)
                remaining = self.policy.wait_seconds - (time.monotonic() - started)
                if remaining <= 0:
                    await active_store.cancel_waiter(bucket, waiter_id)
                    record["capacity_deadline_exceeded"] = True
                    _bind_trace(trace, record)
                    logger.info("provider_capacity_deadline_exceeded", **record)
                    raise ProviderCapacityUnavailableError(
                        "AI provider capacity is busy. Please retry shortly; no build has started."
                    )
                _bind_trace(trace, {**record, "capacity_queued": True})
                await asyncio.sleep(min(self.policy.poll_seconds, remaining))
        except asyncio.CancelledError:
            store = _memory_fallback if used_fallback else self.store
            await store.cancel_waiter(bucket, waiter_id)
            raise
        except ProviderCapacityUnavailableError:
            raise
        except Exception:
            store = _memory_fallback if used_fallback else self.store
            await store.cancel_waiter(bucket, waiter_id)
            raise

    async def release(self, reservation: CapacityReservation) -> None:
        store = _memory_fallback if reservation.used_fallback else self.store
        snapshot = await store.release(reservation.bucket, reservation.demand)
        logger.info(
            "provider_capacity_released",
            provider_capacity_bucket=reservation.bucket.id,
            capacity_dimension=reservation.bucket.dimension,
            reserved_capacity=snapshot.reserved,
            in_flight_capacity=snapshot.in_flight,
            queue_depth=len(snapshot.waiters),
        )

    async def reconcile(
        self, bucket: CapacityBucket, observation: CapacityObservation
    ) -> None:
        now_ms = int(time.time() * 1000)
        snapshot = await self.store.reconcile(bucket, observation, now_ms)
        logger.info(
            "provider_capacity_reconciled",
            provider_capacity_bucket=bucket.id,
            capacity_dimension=observation.dimension,
            observed_limit=snapshot.observed_limit,
            observed_remaining=snapshot.observed_remaining,
            provider_reset_observed=observation.reset_in_seconds,
            reserved_capacity=snapshot.reserved,
            in_flight_capacity=snapshot.in_flight,
        )


def manager_from_settings(settings: Any) -> ProviderCapacityManager:
    return ProviderCapacityManager(store_from_settings(settings), policy_from_settings(settings))


def _trace_fields(
    bucket: CapacityBucket,
    demand: int,
    decision: ReserveDecision,
    snapshot,
    started: float,
) -> dict[str, Any]:
    reserved_after = snapshot.reserved
    reserved_before = reserved_after - demand if decision.granted else reserved_after
    return {
        "provider_capacity_bucket": bucket.id,
        "capacity_dimension": bucket.dimension,
        "observed_limit": decision.observed_limit,
        "observed_remaining": decision.observed_remaining,
        "safe_capacity": decision.safe_capacity,
        "estimated_request_demand": demand,
        "reserved_capacity": reserved_after,
        "local_reserved_before": reserved_before,
        "local_reserved_after": reserved_after,
        "in_flight_capacity": snapshot.in_flight,
        "queue_depth": decision.queue_depth,
        "capacity_wait_ms": int((time.monotonic() - started) * 1000),
        "capacity_mode": decision.mode,
        "capacity_granted": decision.granted,
        "capacity_queued": decision.queued,
        "capacity_immediate_grant": bool(decision.granted and not decision.queued),
        "capacity_deadline_exceeded": False,
        "provider_reset_at_ms": snapshot.reset_at_ms,
    }


def _bind_trace(trace: dict | None, record: Mapping[str, Any]) -> None:
    if trace is None:
        return
    current = dict(trace.get("provider_capacity") or {})
    current.update(record)
    trace["provider_capacity"] = current