import requests, json, time

data = {
    "stream": {
        # "url": "https://www.youtube.com/watch?v=9bZkp7q19f0",
        "url": "https://tagesschau.akamaized.net/hls/live/2020117/tagesschau/tagesschau_3/master-720p-3200.m3u8",
        "media_id": "9bZkp7q19f0",
    }
}
response = requests.post(
    "http://localhost:8000/stream/caption_stream",
    data={"body": json.dumps(data)},
    stream=True,
)
init_timestamp = None
init_time = None
for chunk in response.iter_content(chunk_size=None):
    chunk = json.loads(chunk)

    if not init_timestamp:
        init_timestamp = chunk["timestamps"][0]/100000
        init_time = time.time()

    # calculate lag
    current_time = time.time()
    time_diff = current_time - init_time
    timestamp_diff = chunk["timestamps"][0]/100000 - init_timestamp
    lag = time_diff - timestamp_diff
    print(f"Time diff: {time_diff:.2f}s")
    print(f"Timestamp diff: {timestamp_diff:.2f}s")
    print(f"Lag: {lag:.2f}s")

    print (chunk)

async def run_async_generator_concurrently(async_generator, process, batch_size=2):
    queue = asyncio.Queue(batch_size * 2)
    result_queue = asyncio.Queue()
    num_consumers = batch_size

    async def producer():
        async for i in async_generator:
            await queue.put(i)
        # Signal all consumers to shut down by putting a sentinel for each consumer in the queue
        for _ in range(num_consumers):
            await queue.put(None)

    async def consumer():
        while True:
            item = await queue.get()
            if item is None:  # Check for the sentinel
                queue.task_done()
                break
            result = await process(item)
            await result_queue.put(result)
            queue.task_done()

    consumers = [consumer() for _ in range(num_consumers)]
    # # Setup the producer and consumers to run concurrently
    # await asyncio.gather(producer(), *consumers)

    producer_task = asyncio.create_task(producer())
    consumer_tasks = [asyncio.create_task(c) for c in consumers]

    # Yield all results as they are processed
    while True:
        if (
            result_queue.empty()
            and all(c.done() for c in consumer_tasks)
            and producer_task.done()
        ):
            break
        result = await result_queue.get()
        yield result
        result_queue.task_done()

    # Wait for all tasks to complete
    await producer_task
    for task in consumer_tasks:
        await task

import asyncio


async def async_generator(n):
    for i in range(n):
        yield i


async def process(i):
    await asyncio.sleep(0.1)
    print(f"i * i = {i * i}")
    return i * i

gen = async_generator(10)
async for item in run_async_generator_concurrently(gen, process, 5):
    print(item)

class CaptionStreamEndpoint(Endpoint):
    """Transcribe video in chunks endpoint."""

    async def initialize(self):
        """Initialize the endpoint."""
        self.captioning_handle = await AanaDeploymentHandle.create(
            "captioning_deployment"
        )

    async def run(
        self,
        stream: StreamInput,
        batch_size: int,
    ) -> AsyncGenerator[CaptionStreamOutput, None]:
        """Transcribe video in chunks."""
        async def predict_captions(frames_dict):
            captioning_output = await self.captioning_handle.generate_batch(
                images=frames_dict["frames"]
            )
            # print("captioning_output", captioning_output)
            return {
                "captions": captioning_output["captions"],
                "timestamps": frames_dict["timestamps"],
            }

        gen = run_remote(fetch_stream_frames)(stream_input=stream, batch_size=2)
        async for item in run_async_generator_concurrently(
            gen, predict_captions, batch_size
        ):
            yield item