"""CMP opt-in: one continuous chapter worker per GPU, with immediate progress."""
import asyncio
import json
import logging
import httpx

logger = logging.getLogger(__name__)
_lock = asyncio.Lock()


async def warm_chapters(script, count, endpoints, params, disconnected=None):
    if count <= 0:
        return
    if not endpoints:
        raise ValueError('No CMP GPU endpoints configured')
    async with _lock:
        if disconnected and await disconnected():
            raise asyncio.CancelledError()
        async with httpx.AsyncClient(timeout=httpx.Timeout(7200, connect=15), trust_env=False) as client:
            indices = iter(range(count))
            results = asyncio.Queue()

            async def worker(endpoint):
                try:
                    for index in indices:
                        payload = dict(params, text=script, chapter_index=index)
                        response = await client.post(endpoint + '/audiobook/preview', json=payload)
                        response.raise_for_status()
                        result = response.json()
                        logger.info('CMP_GPU_CHAPTER gpu=%s chapter=%s cached=%s duration=%s',
                                    endpoint, index, result.get('cached'), result.get('duration_s'))
                        event = dict(type='chapter', index=index, total=count,
                                     title=result.get('title', ''), duration_s=result['duration_s'],
                                     cached=result.get('cached', False), stage='gpu_render')
                        await results.put(('chapter', event))
                except Exception as error:
                    await results.put(('error', error))
                finally:
                    await results.put(('finished', None))

            tasks = [asyncio.create_task(worker(endpoint)) for endpoint in endpoints[:count]]
            finished = 0
            try:
                while finished < len(tasks):
                    if disconnected and await disconnected():
                        raise asyncio.CancelledError()
                    try:
                        kind, value = await asyncio.wait_for(results.get(), timeout=5)
                    except asyncio.TimeoutError:
                        yield ': cmp GPU rendering\n\n'
                        continue
                    if kind == 'error':
                        raise value
                    if kind == 'finished':
                        finished += 1
                    else:
                        yield 'data: ' + json.dumps(value) + '\n\n'
            finally:
                for task in tasks:
                    if not task.done():
                        task.cancel()
                await asyncio.gather(*tasks, return_exceptions=True)