#!/usr/bin/env python3
"""Regression tests for bounded lease-broker read/handle/send deadlines."""

from __future__ import annotations

import importlib.util
import json
import socket
import threading
import time
import unittest
from pathlib import Path


DAEMON_PATH = Path(__file__).parents[2] / "framework/tools/lease-broker/daemon.py"
SPEC = importlib.util.spec_from_file_location("lease_broker_deadline_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)


def original_connection_budget() -> float:
    read_budget = getattr(DAEMON, "READ_DEADLINE_SECONDS", None)
    if isinstance(read_budget, (int, float)):
        return float(read_budget)
    return float(DAEMON.CONNECTION_DEADLINE_SECONDS)


class SlowBroker:
    def __init__(self, delay: float) -> None:
        self.delay = delay
        self.calls = 0

    def handle(self, _peer: tuple[int, int, int], request: dict[str, object]) -> dict[str, object]:
        self.calls += 1
        time.sleep(self.delay)
        return {"ok": True, "echo": request.get("action")}


class BrokerDeadlineTest(unittest.TestCase):
    def test_lock_queue_timeout_returns_explicit_fail_closed_reply_without_handling(self) -> None:
        broker = SlowBroker(0)
        broker_lock = threading.Lock()
        broker_lock.acquire()
        server, client = socket.socketpair(socket.AF_UNIX, socket.SOCK_STREAM)
        client.settimeout(DAEMON.HANDLE_QUEUE_TIMEOUT_SECONDS + 2.0)
        worker = threading.Thread(
            target=DAEMON.handle_connection,
            args=(server, broker, broker_lock),
            daemon=True,
        )
        worker.start()
        client.sendall(b'{"action":"probe"}\n')
        client.shutdown(socket.SHUT_WR)

        reply = bytearray()
        try:
            while True:
                chunk = client.recv(4096)
                if not chunk:
                    break
                reply.extend(chunk)
        finally:
            broker_lock.release()
            worker.join(timeout=2.0)
            client.close()

        self.assertFalse(worker.is_alive())
        self.assertEqual(broker.calls, 0)
        self.assertEqual(json.loads(reply), {"ok": False, "code": "BROKER_BUSY"})

    def test_completed_slow_handle_gets_a_complete_framed_reply(self) -> None:
        broker = SlowBroker(original_connection_budget() + 0.1)
        broker_lock = threading.Lock()
        server, client = socket.socketpair(socket.AF_UNIX, socket.SOCK_STREAM)
        client.settimeout(original_connection_budget() + 2.0)
        worker = threading.Thread(
            target=DAEMON.handle_connection,
            args=(server, broker, broker_lock),
            daemon=True,
        )
        worker.start()
        client.sendall(b'{"action":"probe"}\n')
        client.shutdown(socket.SHUT_WR)

        reply = bytearray()
        while True:
            chunk = client.recv(4096)
            if not chunk:
                break
            reply.extend(chunk)
        worker.join(timeout=2.0)
        client.close()

        self.assertFalse(worker.is_alive())
        self.assertEqual(broker.calls, 1)
        self.assertTrue(reply.endswith(b"\n"), f"unframed reply: {bytes(reply)!r}")
        self.assertEqual(json.loads(reply), {"ok": True, "echo": "probe"})


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