from __future__ import annotations

from pathlib import Path
from datetime import datetime, timezone
import tempfile
import unittest

from hibor_fast_collection.budget import Budget, BudgetExpired
from hibor_fast_collection.models import (
    ContractError, ErrorCode, EvidenceState, PackageState, TaskSpec, TerminalReceipt,
)
from helpers import task_spec


class FakeClock:
    def __init__(self): self.value = 0.0
    def __call__(self): return self.value


class ModelsBudgetTests(unittest.TestCase):
    def test_task_spec_exact_and_batch_budget(self):
        with tempfile.TemporaryDirectory() as tmp:
            spec = TaskSpec.from_mapping(task_spec(Path(tmp), quantity=3))
            self.assertEqual(spec.total_budget_ms, 1_080_000)
            bad = task_spec(Path(tmp))
            bad["package_name"] = "other"
            with self.assertRaises(ContractError): TaskSpec.from_mapping(bad)

    def test_key_order_rejected(self):
        with tempfile.TemporaryDirectory() as tmp:
            value = task_spec(Path(tmp))
            first = value.pop("schema_version")
            value["schema_version"] = first
            with self.assertRaises(ContractError): TaskSpec.from_mapping(value)

    def test_budget_work_and_close_deadline(self):
        clock = FakeClock()
        budget = Budget(10_000, 2_000, clock=clock)
        self.assertEqual(budget.require("x", "y", 9_000), 8_000)
        clock.value = 8.001
        with self.assertRaises(BudgetExpired): budget.require("x", "y", 1_000)
        self.assertGreater(budget.require("close", "z", 1_000, close=True), 0)
        clock.value = 10.1
        with self.assertRaises(BudgetExpired): budget.require("close", "z", 1_000, close=True)

    def test_observed_deadline_and_batch_increment(self):
        fresh_clock = FakeClock()
        fresh = Budget(840_000, 30_000, batch_increment_ms=240_000,
                       clock=fresh_clock)
        fresh.begin_item(1)
        self.assertEqual((fresh.remaining_ms(), fresh.remaining_ms(close=True),
                          fresh.remaining_ms(batch_close=True)),
                         (570_000, 600_000, 840_000))

        clock = FakeClock()
        wall = lambda: datetime(2026, 7, 29, 0, 5, tzinfo=timezone.utc)
        budget = Budget(840_000, 30_000, observed_at_utc="2026-07-29T00:00:00Z",
                        batch_increment_ms=240_000, clock=clock, utc_now=wall)
        budget.begin_item(1)
        self.assertEqual(budget.remaining_ms(), 270_000)
        self.assertEqual(budget.remaining_ms(close=True), 300_000)
        self.assertEqual(budget.remaining_ms(batch_close=True), 540_000)
        with self.assertRaises(ValueError):
            budget.begin_item(2)
        clock.value = 299.0
        self.assertTrue(budget.complete_item(1))
        budget.begin_item(2)
        self.assertEqual(budget.remaining_ms(), 210_000)
        self.assertEqual(budget.remaining_ms(close=True), 240_000)
        clock.value = 509.1
        with self.assertRaises(BudgetExpired): budget.checkpoint("item_2", "terminal")
        self.assertGreater(budget.remaining_ms(close=True), 0)
        clock.value = 539.1
        self.assertFalse(budget.complete_item(2))

    def test_terminal_receipt_cross_fields(self):
        target = Path("terminal.json")
        valid = TerminalReceipt(target, target, True, True, EvidenceState.V, True, False, True,
                                10, "a"*64, None, PackageState.P09_TIMING)
        valid.validate()
        file_u = TerminalReceipt(target, None, True, None, EvidenceState.U, True, False, True,
                                 None, None, ErrorCode.RECOVERY_UNKNOWN, PackageState.P09_TIMING)
        file_u.validate()
        live_u = TerminalReceipt(target, None, True, None, EvidenceState.U, None, True, True,
                                 None, None, ErrorCode.PROCESS_LIVENESS_UNKNOWN, PackageState.P09_TIMING)
        live_u.validate()
        with self.assertRaises(ContractError):
            TerminalReceipt(target, None, False, False, EvidenceState.U, False, False, True,
                            None, None, ErrorCode.RECOVERY_UNKNOWN, PackageState.P09_TIMING).validate()


if __name__ == "__main__": unittest.main()