from __future__ import annotations

import threading
import time
import unittest

from memnixa.gateway.run_queue import InMemoryRunQueue


class InMemoryRunQueueTests(unittest.TestCase):
    def test_same_session_jobs_run_serially(self) -> None:
        queue = InMemoryRunQueue(max_concurrent=2)
        first_started = threading.Event()
        second_started = threading.Event()
        release_first = threading.Event()
        order: list[str] = []

        def first_job() -> str:
            order.append("start-1")
            first_started.set()
            release_first.wait(timeout=1)
            order.append("end-1")
            return "one"

        def second_job() -> str:
            order.append("start-2")
            second_started.set()
            order.append("end-2")
            return "two"

        future_one = queue.enqueue("session-a", first_job)
        self.assertTrue(first_started.wait(timeout=1))
        future_two = queue.enqueue("session-a", second_job)
        time.sleep(0.05)
        self.assertFalse(second_started.is_set())
        release_first.set()

        self.assertEqual(future_one.result(timeout=1), "one")
        self.assertEqual(future_two.result(timeout=1), "two")
        self.assertEqual(order, ["start-1", "end-1", "start-2", "end-2"])
        queue.shutdown(wait=False)

    def test_global_concurrency_limit_applies_across_sessions(self) -> None:
        queue = InMemoryRunQueue(max_concurrent=1)
        first_started = threading.Event()
        second_started = threading.Event()
        release_first = threading.Event()

        def first_job() -> str:
            first_started.set()
            release_first.wait(timeout=1)
            return "one"

        def second_job() -> str:
            second_started.set()
            return "two"

        future_one = queue.enqueue("session-a", first_job)
        self.assertTrue(first_started.wait(timeout=1))
        future_two = queue.enqueue("session-b", second_job)
        time.sleep(0.05)
        self.assertFalse(second_started.is_set())
        release_first.set()

        self.assertEqual(future_one.result(timeout=1), "one")
        self.assertEqual(future_two.result(timeout=1), "two")
        snapshot = queue.snapshot()
        self.assertEqual(snapshot["active_count"], 0)
        self.assertEqual(snapshot["queued_count"], 0)
        queue.shutdown(wait=False)


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