"""In-memory queue service for Copilot proxy requests."""
from __future__ import annotations

import asyncio
import logging
from typing import Any

from prompt_auto_lab.api.models import CopilotRequest, CopilotResponse, CopilotModelInfo

logger = logging.getLogger(__name__)


class CopilotQueueService:
    """Manages pending LLM requests for Copilot proxy.

    Thread-safe in-memory queue that:
    - Accepts requests from CopilotProxyClient
    - Serves pending requests to VS Code extension via polling
    - Receives responses and unblocks waiting clients
    """

    def __init__(self) -> None:
        self._pending: dict[str, CopilotRequest] = {}
        self._events: dict[str, asyncio.Event] = {}
        self._results: dict[str, CopilotResponse] = {}
        self._available_models: list[CopilotModelInfo] = []
        self._lock = asyncio.Lock()

    def is_extension_connected(self) -> bool:
        """Check if VS Code extension has registered models."""
        return len(self._available_models) > 0

    def get_available_models(self) -> list[CopilotModelInfo]:
        """Get list of models registered by VS Code extension."""
        return self._available_models.copy()

    async def set_available_models(self, models: list[CopilotModelInfo]) -> None:
        """Update available models (called by extension on connect)."""
        async with self._lock:
            self._available_models = models
            logger.info(f"Copilot models updated: {[m.id for m in models]}")

    async def clear_models(self) -> None:
        """Clear models (called when extension disconnects)."""
        async with self._lock:
            self._available_models = []
            logger.info("Copilot models cleared (extension disconnected)")

    async def queue_request(self, request: CopilotRequest) -> str:
        """Add a request to the queue. Returns request_id."""
        async with self._lock:
            self._pending[request.request_id] = request
            self._events[request.request_id] = asyncio.Event()
            logger.debug(f"Queued Copilot request: {request.request_id}")
        return request.request_id

    async def get_pending(self) -> list[CopilotRequest]:
        """Get all pending requests (for extension polling)."""
        async with self._lock:
            return list(self._pending.values())

    async def submit_response(self, request_id: str, response: CopilotResponse) -> bool:
        """Submit a response for a pending request. Returns True if request existed."""
        async with self._lock:
            if request_id not in self._pending:
                logger.warning(f"Response for unknown request: {request_id}")
                return False

            del self._pending[request_id]
            self._results[request_id] = response

            if request_id in self._events:
                self._events[request_id].set()

            logger.debug(f"Received Copilot response: {request_id}")
            return True

    async def wait_for_response(
        self, request_id: str, timeout: float = 300.0
    ) -> CopilotResponse:
        """Wait for response to a queued request. Raises TimeoutError if timeout."""
        event = self._events.get(request_id)
        if event is None:
            raise KeyError(f"Unknown request_id: {request_id}")

        try:
            await asyncio.wait_for(event.wait(), timeout=timeout)
        except asyncio.TimeoutError:
            # Clean up on timeout
            async with self._lock:
                self._pending.pop(request_id, None)
                self._events.pop(request_id, None)
                self._results.pop(request_id, None)
            raise TimeoutError(f"Copilot request timed out after {timeout}s: {request_id}")

        async with self._lock:
            self._events.pop(request_id, None)
            result = self._results.pop(request_id)

        return result

    async def cancel_request(self, request_id: str) -> bool:
        """Cancel a pending request. Returns True if request existed."""
        async with self._lock:
            if request_id in self._pending:
                del self._pending[request_id]
                if request_id in self._events:
                    self._events[request_id].set()  # Unblock waiter
                return True
            return False


# Global singleton instance
_copilot_queue: CopilotQueueService | None = None


def get_copilot_queue() -> CopilotQueueService:
    """Get the global Copilot queue service instance."""
    global _copilot_queue
    if _copilot_queue is None:
        _copilot_queue = CopilotQueueService()
    return _copilot_queue
