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

spec = importlib.util.spec_from_file_location(
    'gpu_hook', Path(__file__).resolve().parents[1] / 'pve-shared-gpu-hook.py')
hook = importlib.util.module_from_spec(spec)
spec.loader.exec_module(hook)


class SharedGPUHookTests(unittest.TestCase):
    def setUp(self):
        self.directory = tempfile.TemporaryDirectory()
        self.addCleanup(self.directory.cleanup)
        root = Path(self.directory.name)
        (root / 'boot').write_text('test-boot')
        for name, value in [('STATE_DIR', root), ('BOOT_ID', root / 'boot')]:
            p = patch.object(hook, name, value)
            p.start()
            self.addCleanup(p.stop)

    def test_running_guest_or_busy_device_never_resets(self):
        for failing_guard in ('stopped', 'unused_gpu'):
            with self.subTest(guard=failing_guard), patch.object(hook, 'stopped'), \
                    patch.object(hook, 'unused_gpu'), patch.object(hook, 'reenumerate') as reset:
                getattr(hook, failing_guard).side_effect = RuntimeError('busy')
                with self.assertRaises(RuntimeError):
                    hook.handle(102, 'pre-start')
                reset.assert_not_called()
                self.assertFalse((hook.STATE_DIR / 'reservation.json').exists())

    def test_second_start_blocked_until_first_guest_stops(self):
        with patch.object(hook, 'stopped'), patch.object(hook, 'unused_gpu'), \
                patch.object(hook, 'reenumerate') as reset:
            hook.handle(102, 'pre-start')
            with self.assertRaises(RuntimeError):
                hook.handle(100, 'pre-start')
            self.assertEqual(reset.call_count, 1)
            hook.handle(100, 'post-stop')
            self.assertTrue((hook.STATE_DIR / 'reservation.json').exists())
            hook.handle(102, 'post-start')
            hook.handle(102, 'post-stop')
            hook.handle(100, 'pre-start')
            self.assertEqual(reset.call_count, 2)

    def test_failed_reset_stays_reserved_and_is_not_replayed(self):
        with patch.object(hook, 'stopped'), patch.object(hook, 'unused_gpu'), \
                patch.object(hook, 'reenumerate', side_effect=RuntimeError('device missing')) as reset:
            with self.assertRaises(RuntimeError):
                hook.handle(102, 'pre-start')
            hook.handle(102, 'post-stop')
            state = json.loads((hook.STATE_DIR / 'reservation.json').read_text())
            self.assertEqual(state['phase'], 'failed')
            with self.assertRaises(RuntimeError):
                hook.handle(102, 'pre-start')
            self.assertEqual(reset.call_count, 1)

    def test_post_start_cannot_claim_another_guest_reservation(self):
        with patch.object(hook, 'stopped'), patch.object(hook, 'unused_gpu'), \
                patch.object(hook, 'reenumerate'):
            hook.handle(102, 'pre-start')
            with self.assertRaises(RuntimeError):
                hook.handle(100, 'post-start')


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