import asyncio
import json
import sys
import unittest
from pathlib import Path
from unittest.mock import patch

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'backend'))
from services.cmp_multigpu import warm_chapters


class DynamicQueueTests(unittest.IsolatedAsyncioTestCase):
    async def test_free_gpu_continues_and_reports_before_slow_gpu_finishes(self):
        slow = asyncio.Event()
        dispatched = []
        active = set()

        class Response:
            def raise_for_status(self): pass
            def json(self): return {'duration_s': 1, 'cached': False}

        class Client:
            def __init__(self, **kwargs): pass
            async def __aenter__(self): return self
            async def __aexit__(self, *args): pass
            async def post(self, url, json):
                self_outer.assertNotIn(url, active)
                active.add(url)
                index = json['chapter_index']
                dispatched.append((url, index))
                try:
                    if index == 0:
                        await slow.wait()
                    else:
                        await asyncio.sleep(.001)
                    return Response()
                finally:
                    active.remove(url)

        self_outer = self
        with patch('services.cmp_multigpu.httpx.AsyncClient', Client):
            stream = warm_chapters('text', 6, ['gpu0', 'gpu1', 'gpu2'], {})
            try:
                first = await asyncio.wait_for(anext(stream), .3)
                self.assertTrue(first.startswith('data:'))
                self.assertFalse(slow.is_set())
                while len(dispatched) < 6:
                    await asyncio.wait_for(anext(stream), .3)
                self.assertTrue(any(index >= 3 and url != 'gpu0/audiobook/preview' for url, index in dispatched))
                slow.set()
                remaining = [event async for event in stream]
                self.assertTrue(any('"index": 0' in event for event in remaining))
                self.assertEqual(list(range(6)), sorted(index for _, index in dispatched))
            finally:
                slow.set()
                await stream.aclose()
        self.assertFalse(active)

    async def test_closing_stream_cancels_inflight_requests(self):
        stopped = asyncio.Event()
        class Response:
            def raise_for_status(self): pass
            def json(self): return {'duration_s': 1}
        class Client:
            def __init__(self, **kwargs): pass
            async def __aenter__(self): return self
            async def __aexit__(self, *args): pass
            async def post(self, url, json):
                if json['chapter_index'] == 0:
                    try: await asyncio.Event().wait()
                    finally: stopped.set()
                return Response()
        with patch('services.cmp_multigpu.httpx.AsyncClient', Client):
            stream = warm_chapters('text', 2, ['gpu0', 'gpu1'], {})
            try:
                await asyncio.wait_for(anext(stream), .3)
            finally:
                await stream.aclose()
            self.assertTrue(stopped.is_set())


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