"""Client runner manager for async request execution."""

import asyncio
import threading
import time
from queue import Empty, Queue
from typing import List, Optional

from veeksha.client.base import BaseLLMClient
from veeksha.core.response import RequestResult
from veeksha.logger import init_logger

logger = init_logger(__name__)


QUEUE_GET_TIMEOUT_S = 0.1


class ClientWorker:
    """Async client worker that processes requests from an input queue."""

    def __init__(
        self,
        worker_id: int,
        client: BaseLLMClient,
        input_queue: Queue,
        output_queue: Queue,
        stop_event: threading.Event,
    ):
        self.worker_id = worker_id
        self.client = client
        self.input_queue = input_queue
        self.output_queue = output_queue
        self.stop_event = stop_event

    def run(self) -> None:
        """Run the async event loop for this worker."""
        asyncio.run(self._run_async())

    async def _run_async(self) -> None:
        """Async main loop."""
        loop = asyncio.get_running_loop()
        active_tasks = set()

        while not self.stop_event.is_set():
            try:
                # avoid blocking the event loop
                item = await loop.run_in_executor(
                    None, lambda: self.input_queue.get(timeout=QUEUE_GET_TIMEOUT_S)
                )
            except Empty:
                continue
            except Exception:
                if self.stop_event.is_set():
                    break
                continue

            if item is None:  # sentinel
                break

            task = asyncio.create_task(self._process_request(item))
            active_tasks.add(task)
            task.add_done_callback(active_tasks.discard)

        if active_tasks:
            logger.debug(
                "Client worker %d cancelling %d pending tasks",
                self.worker_id,
                len(active_tasks),
            )
            for task in active_tasks:
                task.cancel()
            await asyncio.wait(active_tasks, timeout=2.0)

        logger.debug("Client worker %d exiting", self.worker_id)

    async def _process_request(self, item) -> None:
        """Process a single request item."""
        (
            request,
            session_id,
            session_size,
            scheduler_ready_at,
            scheduler_dispatched_at,
        ) = item

        client_picked_up_at: float = time.monotonic()

        try:
            result = await self.client.send_request(
                request=request,
                session_id=session_id,
                session_total_requests=session_size,
            )
        except Exception as e:
            logger.exception(
                f"Client worker {self.worker_id}: Client raised unhandled exception"
            )

            result = RequestResult(
                request_id=request.id,
                session_id=session_id,
                session_total_requests=session_size,
                channels={},
                success=False,
                error_code=500,
                error_msg=f"Unhandled client exception: {str(e)}",
                client_completed_at=time.monotonic(),
            )

        # Set lifecycle timestamps
        result.scheduler_ready_at = scheduler_ready_at
        result.scheduler_dispatched_at = scheduler_dispatched_at
        result.client_picked_up_at = client_picked_up_at

        self.output_queue.put(result)


class ClientRunnerManager:
    """Manager for a pool of async client worker threads."""

    def __init__(
        self,
        client: BaseLLMClient,
        input_queues: List[Queue],
        output_queue: Queue,
        stop_event: threading.Event,
    ):
        """Initialize the client runner manager.

        Args:
            client: LLM client to use for requests
            input_queues: One input queue per worker
            output_queue: Shared output queue for results
            stop_event: Stop event for graceful shutdown
        """
        self.client = client
        self.input_queues = input_queues
        self.output_queue = output_queue
        self.stop_event = stop_event
        self.workers: List[ClientWorker] = []
        self.threads: List[threading.Thread] = []

    def start(self) -> None:
        """Start all client worker threads."""
        for i, queue in enumerate(self.input_queues):
            worker = ClientWorker(
                worker_id=i,
                client=self.client,
                input_queue=queue,
                output_queue=self.output_queue,
                stop_event=self.stop_event,
            )
            self.workers.append(worker)

            thread = threading.Thread(
                target=worker.run,
                name=f"client-worker-{i}",
                daemon=True,
            )
            self.threads.append(thread)
            thread.start()

        logger.info("Started %d client worker threads", len(self.threads))

    def stop(self) -> None:
        """Signal workers to stop."""
        self.stop_event.set()
        for queue in self.input_queues:
            queue.put(None)  # Sentinel

    def wait(self, timeout: Optional[float] = None) -> bool:
        """Wait for all worker threads to finish.

        Args:
            timeout: Maximum seconds to wait for all workers. ``None`` for
                indefinitely.

        Returns:
            True if every worker exited in time, False if the timeout expired.
        """
        start_time = time.monotonic()
        for thread in self.threads:
            join_timeout = None
            if timeout is not None:
                elapsed = time.monotonic() - start_time
                remaining = max(0.0, timeout - elapsed)
                if remaining == 0.0 and thread.is_alive():
                    break
                join_timeout = remaining
            thread.join(join_timeout)

        alive_threads = [thread for thread in self.threads if thread.is_alive()]
        if alive_threads:
            logger.warning(
                "Timed out waiting for %d client worker threads", len(alive_threads)
            )
            return False

        logger.debug("All client worker threads joined")
        return True

    def get_worker_count(self) -> int:
        """Return number of workers."""
        return len(self.workers)
