"""Frozen per-run budgets and fail-closed durable enforcement."""

from __future__ import annotations

from datetime import UTC, datetime
from typing import Any

from deerflow.config.report_collaboration_config import ReportCollaborationConfig
from deerflow.persistence.report_collaboration.base import ReportCollaborationLimitError, ReportCollaborationStore
from deerflow.persistence.report_collaboration.codec import parse_lease_until


def freeze_run_budget(config: ReportCollaborationConfig) -> dict[str, Any]:
    return {
        "max_team_members": config.max_team_members,
        "max_parallel_tasks": config.max_parallel_tasks,
        "max_model_calls": config.max_model_calls,
        "max_retrieval_calls": config.max_retrieval_calls,
        "max_source_bytes": config.max_source_bytes,
        "token_budget": config.token_budget,
        "cost_budget": config.cost_budget,
        "session_timeout_seconds": config.session_timeout_seconds,
        "node_timeout_seconds": config.node_timeout_seconds,
        "max_node_attempts": config.max_node_attempts,
        "max_repair_rounds": config.max_repair_rounds,
    }


class RunBudgetEnforcer:
    def __init__(self, store: ReportCollaborationStore, run_id: str) -> None:
        self._store = store
        self.run_id = run_id

    async def consume(self, deltas: dict[str, float | int], *, action: str, resource_id: str | None = None) -> dict[str, Any]:
        try:
            result = await self._store.consume_run_budget(self.run_id, deltas=deltas)
        except ReportCollaborationLimitError as exc:
            await self.fail(exc.code, str(exc), detail={"attempted_usage": exc.usage, "limits": exc.limits})
            raise
        run = await self._store.get_run_record(self.run_id)
        if run is not None:
            await self._store.append_audit(
                session_id=str(run["session_id"]),
                run_id=self.run_id,
                category="budget",
                action=action,
                resource_type="run",
                resource_id=resource_id or self.run_id,
                detail={"delta": dict(deltas), "usage": result["usage"]},
            )
        return result

    async def remaining_seconds(self) -> float | None:
        run = await self._store.get_run_record(self.run_id)
        if run is None:
            return None
        deadline = parse_lease_until(run.get("deadline_at"))
        if deadline is None:
            return None
        return max(0.0, (deadline - datetime.now(UTC)).total_seconds())

    async def assert_time_available(self) -> None:
        remaining = await self.remaining_seconds()
        if remaining is not None and remaining <= 0:
            await self.fail("RUN_TIME_BUDGET_EXCEEDED", "运行总时长已超过预算")
            raise ReportCollaborationLimitError("RUN_TIME_BUDGET_EXCEEDED", "运行总时长已超过预算")

    async def fail(self, code: str, message: str, *, detail: dict[str, Any] | None = None) -> dict[str, Any]:
        before = await self._store.get_run_record(self.run_id)
        if before is None:
            return {}
        if str(before.get("status") or "") in {"completed", "cancelled"}:
            return await self._store.get_run(self.run_id) or {}
        already_failed = bool(before.get("budget_exhausted_reason"))
        run = await self._store.fail_run_for_limit(self.run_id, reason=code)
        if already_failed:
            return run
        session_id = str(before["session_id"])
        await self._store.append_audit(
            session_id=session_id,
            run_id=self.run_id,
            category="budget",
            action="run.limit_exceeded",
            resource_type="run",
            resource_id=self.run_id,
            outcome="denied",
            reason=code,
            detail={"message": message, **(detail or {})},
        )
        await self._store.append_event(
            session_id=session_id,
            run_id=self.run_id,
            event_type="run.error",
            data={"code": code, "message": message},
        )
        await self._store.append_event(
            session_id=session_id,
            run_id=self.run_id,
            event_type="run.status.changed",
            data={"status": "failed", "reason": code},
        )
        return run


__all__ = ["RunBudgetEnforcer", "freeze_run_budget"]