#!/usr/bin/env python3
"""Standard-library edge tests for lease-broker atomic state persistence."""

from __future__ import annotations

import copy
import importlib.util
import json
import os
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch


DAEMON_PATH = Path(__file__).parents[2] / "framework/tools/lease-broker/daemon.py"
SPEC = importlib.util.spec_from_file_location("lease_broker_daemon", DAEMON_PATH)
if SPEC is None or SPEC.loader is None:
    raise RuntimeError("unable to load lease broker daemon")
DAEMON = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(DAEMON)


class StateStoreCommitTest(unittest.TestCase):
    def make_store(self, root: Path):
        os.chmod(root, 0o700)
        store = DAEMON.StateStore(root / "state.json")
        store.value["marker"] = "partial-write-proof"
        return store

    def test_partial_writes_persist_the_complete_payload(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            store = self.make_store(root)
            real_write = os.write

            def partial_write(descriptor: int, payload: bytes) -> int:
                return real_write(descriptor, payload[: max(1, len(payload) // 3)])

            with patch.object(DAEMON.os, "write", side_effect=partial_write):
                store.commit()

            self.assertEqual(json.loads(store.path.read_text()), store.value)

    def test_zero_progress_removes_owned_temporary_file(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            store = self.make_store(root)
            with patch.object(DAEMON.os, "write", return_value=0):
                with self.assertRaises(OSError):
                    store.commit()

            self.assertFalse(store.path.exists())
            self.assertEqual(list(root.glob(".*.tmp")), [])

    def test_oversized_payload_is_refused_before_replacing_state(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            store = self.make_store(root)
            store.value.pop("marker")
            store.commit()
            durable = store.path.read_bytes()
            store.value["oversized"] = "x" * DAEMON.MAX_STATE

            with patch.object(DAEMON.os, "open", wraps=os.open) as mocked_open:
                with self.assertRaisesRegex(DAEMON.BrokerFailure, "STATE_TOO_LARGE"):
                    store.commit()

            self.assertEqual(mocked_open.call_count, 0)
            self.assertEqual(store.path.read_bytes(), durable)
            self.assertEqual(list(root.glob(".*.tmp")), [])


class StateStoreValidationTest(unittest.TestCase):
    @staticmethod
    def binding() -> dict[str, object]:
        return {
            "compaction_epoch": 0,
            "request_epoch": 0,
            "h_source": "a" * 64,
            "h_payload": "b" * 64,
            "schema_version": 1,
        }

    def test_impossible_or_over_capacity_token_state_is_rejected(self) -> None:
        session_id = "1" * 64
        session = {
            "anchor_pid": 123,
            "anchor_starttime": "456",
            "runtime_generation": 2,
        }
        live_token = {
            "session_id": session_id,
            "runtime_generation": 2,
            "binding": self.binding(),
            "consumed": False,
        }
        cases = {
            "stale generation": {
                "2" * 64: {**live_token, "runtime_generation": 1},
            },
            "consumed token": {
                "2" * 64: {**live_token, "consumed": True},
            },
            "over capacity": {
                f"{index:064x}": copy.deepcopy(live_token)
                for index in range(DAEMON.MAX_PENDING_TOKENS + 1)
            },
        }

        for label, tokens in cases.items():
            with self.subTest(label=label), tempfile.TemporaryDirectory() as directory:
                root = Path(directory)
                os.chmod(root, 0o700)
                state_path = root / "state.json"
                state_path.write_text(json.dumps({
                    "version": 1,
                    "sessions": {session_id: session},
                    "tokens": tokens,
                }))
                os.chmod(state_path, 0o600)

                with self.assertRaisesRegex(DAEMON.BrokerFailure, "STATE_INTEGRITY"):
                    DAEMON.StateStore(state_path)


class BrokerBehaviorTest(unittest.TestCase):
    def make_broker(self, root: Path):
        os.chmod(root, 0o700)
        return DAEMON.Broker(DAEMON.StateStore(root / "state.json"))

    @staticmethod
    def binding() -> dict[str, object]:
        return {
            "compaction_epoch": 0,
            "request_epoch": 0,
            "h_source": "a" * 64,
            "h_payload": "b" * 64,
            "schema_version": 1,
        }

    def register(self, broker, generation: int = 1) -> str:
        response = broker.handle((123, 1000, 1000), {
            "action": "register_anchor",
            "runtime_generation": generation,
        })
        return response["session_id"]

    def mint(self, broker, session_id: str, generation: int = 1) -> str:
        response = broker.handle((123, 1000, 1000), {
            "action": "mint_token",
            "session_id": session_id,
            "runtime_generation": generation,
            "binding": self.binding(),
        })
        return response["token"]

    def consume(self, broker, session_id: str, token: str, generation: int = 1):
        return broker.handle((123, 1000, 1000), {
            "action": "consume_token",
            "session_id": session_id,
            "runtime_generation": generation,
            "token": token,
        })

    def test_anchor_generation_bump_reuses_session_and_revokes_token(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            broker = self.make_broker(root)
            with (
                patch.object(
                    DAEMON,
                    "proc_node",
                    return_value={"pid": 123, "ppid": 1, "starttime": "456"},
                ),
                patch.object(DAEMON, "verified_ancestry", return_value=True),
            ):
                first = broker.handle((123, 1000, 1000), {
                    "action": "register_anchor",
                    "runtime_generation": 1,
                })
                minted = broker.handle((123, 1000, 1000), {
                    "action": "mint_token",
                    "session_id": first["session_id"],
                    "runtime_generation": 1,
                    "binding": self.binding(),
                })
                bumped = broker.handle((123, 1000, 1000), {
                    "action": "register_anchor",
                    "runtime_generation": 2,
                })
                repeated = broker.handle((123, 1000, 1000), {
                    "action": "register_anchor",
                    "runtime_generation": 2,
                })

            self.assertEqual(bumped["session_id"], first["session_id"])
            self.assertEqual(repeated["session_id"], first["session_id"])
            self.assertNotIn(minted["token"], broker.store.tokens())
            restarted = self.make_broker(root)
            self.assertEqual(restarted.store.tokens(), {})
            self.assertEqual(
                restarted.store.sessions()[first["session_id"]]["runtime_generation"], 2
            )
            with self.assertRaisesRegex(DAEMON.BrokerFailure, "STALE_GENERATION"):
                with patch.object(DAEMON, "proc_node", return_value={"starttime": "456"}):
                    broker.handle((123, 1000, 1000), {
                        "action": "register_anchor",
                        "runtime_generation": 1,
                    })

    def test_successful_consume_deletes_token_and_replay_is_refused(self) -> None:
        with tempfile.TemporaryDirectory() as directory, (
            patch.object(DAEMON, "proc_node", return_value={"pid": 123, "ppid": 1, "starttime": "456"})
        ), patch.object(DAEMON, "verified_ancestry", return_value=True):
            broker = self.make_broker(Path(directory))
            session_id = self.register(broker)
            token = self.mint(broker, session_id)

            self.assertEqual(self.consume(broker, session_id, token), {"ok": True})
            self.assertNotIn(token, broker.store.tokens())
            with self.assertRaisesRegex(DAEMON.BrokerFailure, "TOKEN_REPLAY"):
                self.consume(broker, session_id, token)

    def test_normal_cycles_remain_bounded_and_restartable(self) -> None:
        with tempfile.TemporaryDirectory() as directory, (
            patch.object(DAEMON, "proc_node", return_value={"pid": 123, "ppid": 1, "starttime": "456"})
        ), patch.object(DAEMON, "verified_ancestry", return_value=True):
            root = Path(directory)
            broker = self.make_broker(root)
            session_id = self.register(broker)

            for _ in range(DAEMON.MAX_PENDING_TOKENS * 3):
                self.consume(broker, session_id, self.mint(broker, session_id))

            self.assertEqual(broker.store.tokens(), {})
            self.assertLess((root / "state.json").stat().st_size, DAEMON.MAX_STATE)
            restarted = self.make_broker(root)
            self.assertEqual(restarted.store.tokens(), {})
            self.assertIn(session_id, restarted.store.sessions())

    def test_pending_token_capacity_refusal_does_not_mutate_state(self) -> None:
        with tempfile.TemporaryDirectory() as directory, (
            patch.object(DAEMON, "proc_node", return_value={"pid": 123, "ppid": 1, "starttime": "456"})
        ), patch.object(DAEMON, "verified_ancestry", return_value=True):
            root = Path(directory)
            broker = self.make_broker(root)
            session_id = self.register(broker)
            for _ in range(DAEMON.MAX_PENDING_TOKENS):
                self.mint(broker, session_id)
            before = copy.deepcopy(broker.store.value)
            durable = broker.store.path.read_bytes()

            with self.assertRaisesRegex(DAEMON.BrokerFailure, "TOKEN_CAPACITY"):
                self.mint(broker, session_id)

            self.assertEqual(broker.store.value, before)
            self.assertEqual(broker.store.path.read_bytes(), durable)

    def test_directory_fsync_failure_poisoned_store_cannot_continue(self) -> None:
        with tempfile.TemporaryDirectory() as directory, patch.object(
            DAEMON,
            "proc_node",
            return_value={"pid": 123, "ppid": 1, "starttime": "456"},
        ):
            root = Path(directory)
            broker = self.make_broker(root)
            real_fsync = os.fsync

            def fail_directory_fsync(descriptor: int) -> None:
                if os.path.isdir(f"/proc/self/fd/{descriptor}"):
                    raise OSError("directory fsync failed")
                real_fsync(descriptor)

            with patch.object(DAEMON.os, "fsync", side_effect=fail_directory_fsync):
                with self.assertRaisesRegex(
                    DAEMON.StateCommitUncertain, "STATE_COMMIT_UNCERTAIN"
                ):
                    self.register(broker)

            durable = json.loads(broker.store.path.read_text())
            self.assertEqual(broker.store.value, durable)
            self.assertTrue(broker.store.poisoned)
            before = copy.deepcopy(broker.store.value)
            with self.assertRaisesRegex(
                DAEMON.StateCommitUncertain, "STATE_COMMIT_UNCERTAIN"
            ):
                broker.handle((123, 1000, 1000), {
                    "action": "register_anchor",
                    "runtime_generation": 2,
                })
            self.assertEqual(broker.store.value, before)

    def test_commit_failures_before_replace_roll_back_every_broker_mutation(self) -> None:
        with tempfile.TemporaryDirectory() as directory, (
            patch.object(DAEMON, "proc_node", return_value={"pid": 123, "ppid": 1, "starttime": "456"})
        ), patch.object(DAEMON, "verified_ancestry", return_value=True):
            root = Path(directory)
            broker = self.make_broker(root)
            session_id = self.register(broker)
            token = self.mint(broker, session_id)

            def assert_rollback(request: dict[str, object]) -> None:
                before = copy.deepcopy(broker.store.value)
                durable = broker.store.path.read_bytes()
                with patch.object(broker.store, "commit", side_effect=OSError("fsync failed")):
                    with self.assertRaisesRegex(OSError, "fsync failed"):
                        broker.handle((123, 1000, 1000), request)
                self.assertEqual(broker.store.value, before)
                self.assertEqual(broker.store.path.read_bytes(), durable)

            assert_rollback({"action": "register_anchor", "runtime_generation": 2})
            assert_rollback({
                "action": "mint_token", "session_id": session_id,
                "runtime_generation": 1, "binding": self.binding(),
            })
            assert_rollback({
                "action": "consume_token", "session_id": session_id,
                "runtime_generation": 1, "token": token,
            })

            with tempfile.TemporaryDirectory() as second_directory:
                second = self.make_broker(Path(second_directory))
                with patch.object(second.store, "commit", side_effect=OSError("fsync failed")):
                    with self.assertRaisesRegex(OSError, "fsync failed"):
                        self.register(second)
                self.assertEqual(
                    second.store.value, {"version": 1, "sessions": {}, "tokens": {}}
                )
                self.assertFalse(second.store.path.exists())

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