from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
import time
from typing import Callable

from .models import ContractError, ErrorCode


class BudgetExpired(ContractError):
    def __init__(self, phase: str, operation: str, deadline_scope: str):
        super().__init__(ErrorCode.DEADLINE_EXPIRED, operation,
                         f"{phase}:{deadline_scope}")
        self.phase = phase
        self.operation = operation
        self.deadline_scope = deadline_scope


@dataclass
class Budget:
    total_ms: int
    close_reserve_ms: int
    observed_at_utc: str | None = None
    batch_increment_ms: int = 240_000
    clock: Callable[[], float] = time.monotonic
    utc_now: Callable[[], datetime] = lambda: datetime.now(timezone.utc)

    def __post_init__(self) -> None:
        if self.total_ms <= 0 or self.close_reserve_ms < 0 or self.close_reserve_ms >= self.total_ms:
            raise ValueError("invalid budget")
        if self.batch_increment_ms <= 0:
            raise ValueError("invalid batch increment")
        mono_now = self.clock()
        startup_elapsed = 0.0
        if self.observed_at_utc is not None:
            try:
                observed = datetime.fromisoformat(self.observed_at_utc.replace("Z", "+00:00"))
                if observed.tzinfo is None:
                    raise ValueError("timezone required")
                wall_now = self.utc_now()
                if wall_now.tzinfo is None:
                    raise ValueError("UTC clock must be aware")
                startup_elapsed = (wall_now.astimezone(timezone.utc) - observed.astimezone(timezone.utc)).total_seconds()
            except (TypeError, ValueError, OverflowError) as exc:
                raise BudgetExpired("startup", "parse_observed_at", "work") from exc
            if startup_elapsed < -2.0:
                raise BudgetExpired("startup", "observed_at_in_future", "work")
            startup_elapsed = max(0.0, startup_elapsed)
        self._start = mono_now - startup_elapsed
        self._batch_close_deadline = self._start + self.total_ms / 1000
        self._batch_work_deadline = self._batch_close_deadline - self.close_reserve_ms / 1000
        self._item_close_deadline: float | None = None
        self._item_work_deadline: float | None = None
        self._active_item = 0
        self._item_scope_active = False
        self._previous_terminal: float | None = None

    @property
    def elapsed_ms(self) -> int:
        return max(0, round((self.clock() - self._start) * 1000))

    def remaining_ms(self, *, close: bool = False, batch_close: bool = False) -> int:
        if close and batch_close:
            raise ValueError("close scopes are mutually exclusive")
        if batch_close:
            deadline = self._batch_close_deadline
        elif close:
            deadline = (self._item_close_deadline if self._item_scope_active
                        else self._batch_close_deadline)
        else:
            deadline = (self._item_work_deadline if self._item_scope_active
                        else self._batch_work_deadline)
        if deadline is None:
            raise ValueError("deadline unavailable")
        return max(0, int((deadline - self.clock()) * 1000))

    def require(self, phase: str, operation: str, requested_ms: int, *,
                close: bool = False, batch_close: bool = False) -> int:
        remaining = self.remaining_ms(close=close, batch_close=batch_close)
        if remaining <= 0:
            scope = "batch_close" if batch_close else "item_close" if close else "work"
            raise BudgetExpired(phase, operation, scope)
        return min(remaining, max(1, requested_ms))

    def checkpoint(self, phase: str, operation: str, *,
                   close: bool = False, batch_close: bool = False) -> None:
        self.require(phase, operation, 1, close=close, batch_close=batch_close)

    def begin_item(self, index: int) -> None:
        """Bind I01 to observed+600s and each later item to prior item terminal+240s."""
        if index < 1:
            raise ValueError("item order")
        if self._item_scope_active and index == self._active_item:
            return
        expected = 1 if self._active_item == 0 else self._active_item + 1
        if index != expected or (index > 1 and self._previous_terminal is None):
            raise ValueError("previous item terminal required")
        now = self.clock()
        if index == 1:
            candidate = self._start + 600.0
        else:
            candidate = self._previous_terminal + self.batch_increment_ms / 1000
        self._item_close_deadline = min(candidate, self._batch_close_deadline)
        self._item_work_deadline = self._item_close_deadline - self.close_reserve_ms / 1000
        self._active_item = index
        self._item_scope_active = True
        self.checkpoint(f"item_{index}", "begin")

    @property
    def active_item_deadline(self) -> float | None:
        return self._item_close_deadline

    @property
    def active_item_work_deadline(self) -> float | None:
        return self._item_work_deadline

    def complete_item(self, index: int) -> bool:
        """Record the prior terminal anchor and report whether it met item_close."""
        if not self._item_scope_active or index != self._active_item:
            raise ValueError("active item required")
        now = self.clock()
        within_deadline = bool(self._item_close_deadline is not None and
                               now < self._item_close_deadline)
        self._previous_terminal = now
        self._item_scope_active = False
        return within_deadline

    @property
    def work_deadline_reached(self) -> bool:
        deadline = (self._item_work_deadline if self._item_scope_active
                    else self._batch_work_deadline)
        return bool(deadline is not None and self.clock() >= deadline)

    @property
    def close_deadline_reached(self) -> bool:
        deadline = (self._item_close_deadline if self._item_scope_active
                    else self._batch_close_deadline)
        return bool(deadline is not None and self.clock() >= deadline)

    @property
    def batch_close_deadline_reached(self) -> bool:
        return self.clock() >= self._batch_close_deadline