from __future__ import annotations

from datetime import datetime, timezone

from pydantic import BaseModel

from nexus.core.events.bus import EventBus
from nexus.core.events.types import BudgetExhaustedEvent, BudgetWarningEvent
from nexus.resilience.errors import BudgetExhausted, OrchestratorError


class AgentLedger(BaseModel):
    """Per-agent token ledger."""

    agent_id: str
    allocated: int
    used_input: int = 0
    used_output: int = 0
    call_count: int = 0
    last_call_at: datetime | None = None
    compression_triggers: int = 0

    @property
    def total_used(self) -> int:
        return self.used_input + self.used_output

    @property
    def remaining(self) -> int:
        return max(0, self.allocated - self.total_used)

    @property
    def pct_used(self) -> float:
        if self.allocated <= 0:
            return 0.0
        return self.total_used / self.allocated


class TokenAccountant:
    """Tracks token usage and emits budget-related events."""

    def __init__(self, bus: EventBus, warning_threshold: float = 0.80, hard_cutoff: float = 0.98) -> None:
        self.bus = bus
        self.warning_threshold = warning_threshold
        self.hard_cutoff = hard_cutoff
        self._ledgers: dict[str, AgentLedger] = {}

    def register_agent(self, agent_id: str, budget: int) -> None:
        """Register or reset an agent ledger."""
        self._ledgers[agent_id] = AgentLedger(agent_id=agent_id, allocated=max(1, budget))

    async def record_call(self, agent_id: str, input_tokens: int, output_tokens: int) -> AgentLedger:
        """Record token usage and enforce soft/hard thresholds."""
        if agent_id not in self._ledgers:
            raise OrchestratorError(f"Agent {agent_id} not registered in accountant", agent_id=agent_id)

        ledger = self._ledgers[agent_id]
        ledger.used_input += max(0, input_tokens)
        ledger.used_output += max(0, output_tokens)
        ledger.call_count += 1
        ledger.last_call_at = datetime.now(timezone.utc)

        if ledger.pct_used >= self.warning_threshold:
            ledger.compression_triggers += 1
            await self.bus.publish_and_wait(
                BudgetWarningEvent(
                    source_id=agent_id,
                    agent_id=agent_id,
                    used_tokens=ledger.total_used,
                    allocated_tokens=ledger.allocated,
                    pct=ledger.pct_used,
                )
            )

        if ledger.pct_used >= self.hard_cutoff:
            await self.bus.publish_and_wait(
                BudgetExhaustedEvent(
                    source_id=agent_id,
                    agent_id=agent_id,
                    used_tokens=ledger.total_used,
                    allocated_tokens=ledger.allocated,
                )
            )
            raise BudgetExhausted(
                "Token budget exhausted",
                used_tokens=ledger.total_used,
                allocated_tokens=ledger.allocated,
                agent_id=agent_id,
            )

        return ledger

    def get_ledger(self, agent_id: str) -> AgentLedger:
        """Get an agent ledger by id."""
        if agent_id not in self._ledgers:
            raise OrchestratorError(f"Unknown agent ledger: {agent_id}", agent_id=agent_id)
        return self._ledgers[agent_id]

    def get_all_ledgers(self) -> dict[str, AgentLedger]:
        """Get all ledgers."""
        return dict(self._ledgers)

    def session_total(self) -> int:
        """Total tokens used across all agents."""
        return sum(ledger.total_used for ledger in self._ledgers.values())

    def to_report(self) -> str:
        """Render a simple usage report."""
        lines = ["# Token Report", ""]
        for agent_id in sorted(self._ledgers):
            ledger = self._ledgers[agent_id]
            lines.append(
                (
                    f"- {agent_id}: used={ledger.total_used} "
                    f"(in={ledger.used_input}, out={ledger.used_output}) "
                    f"allocated={ledger.allocated} pct={ledger.pct_used:.1%}"
                )
            )
        lines.append("")
        lines.append(f"session_total={self.session_total()}")
        return "\n".join(lines)
