"""Exercise process-wide MHA flag protection across threads and exceptions."""
import threading
import unittest
from pathlib import Path
import sys

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

import torch
from parametric_skills.runtime import inference_context


class InferenceContextTests(unittest.TestCase):
    def test_nested_context_and_exception_restore_previous_flag(self):
        original = torch.backends.mha.get_fastpath_enabled()
        try:
            for previous in [True, False]:
                torch.backends.mha.set_fastpath_enabled(previous)
                with self.assertRaisesRegex(RuntimeError, "test"):
                    with inference_context(torch.device("cpu"), torch.float32):
                        with inference_context(torch.device("cpu"), torch.float32):
                            self.assertFalse(torch.backends.mha.get_fastpath_enabled())
                        self.assertFalse(torch.backends.mha.get_fastpath_enabled())
                        raise RuntimeError("test")
                self.assertEqual(torch.backends.mha.get_fastpath_enabled(), previous)
        finally:
            torch.backends.mha.set_fastpath_enabled(original)

    def test_overlapping_threads_cannot_restore_flag_during_another_forward(self):
        original = torch.backends.mha.get_fastpath_enabled()
        first_entered, second_attempted, second_entered = (threading.Event() for _ in range(3))
        release_first, release_second = threading.Event(), threading.Event()
        errors = []
        def worker(first):
            try:
                if not first:
                    second_attempted.set()
                with inference_context(torch.device("cpu"), torch.float32):
                    (first_entered if first else second_entered).set()
                    if not (release_first if first else release_second).wait(5):
                        raise TimeoutError("test synchronization")
            except BaseException as exc:
                errors.append(exc)
        threads = [threading.Thread(target=worker, args=(first,)) for first in [True, False]]
        try:
            torch.backends.mha.set_fastpath_enabled(True)
            threads[0].start()
            self.assertTrue(first_entered.wait(5))
            threads[1].start()
            self.assertTrue(second_attempted.wait(5))
            overlap = second_entered.wait(0.1)
            release_first.set()
            threads[0].join(5)
            self.assertTrue(second_entered.wait(5))
            self.assertFalse(torch.backends.mha.get_fastpath_enabled())
            self.assertFalse(overlap, "Concurrent contexts must serialize across instances")
        finally:
            release_first.set()
            release_second.set()
            for thread in threads:
                if thread.ident is not None:
                    thread.join(5)
            final = torch.backends.mha.get_fastpath_enabled()
            torch.backends.mha.set_fastpath_enabled(original)
        self.assertFalse(errors, errors)
        self.assertTrue(final)


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