"""Async batch logger for request logs."""
import asyncio
from typing import Optional, List
from dataclasses import dataclass
from ..database import db_connection
from ..config import LOG_BATCH_SIZE, LOG_FLUSH_INTERVAL


@dataclass
class LogEntry:
    """A single log entry."""
    client_key_id: int
    provider_id: int
    model: str
    prompt_tokens: Optional[int] = None
    completion_tokens: Optional[int] = None
    latency_ms: Optional[int] = None
    success: bool = True
    error_message: Optional[str] = None


class AsyncLogger:
    """Asynchronous batch logger for request logs."""

    def __init__(self):
        self._queue: asyncio.Queue[LogEntry] = asyncio.Queue()
        self._task: Optional[asyncio.Task] = None

    async def start(self):
        """Start the background flush task."""
        if self._task is None:
            self._task = asyncio.create_task(self._flush_loop())

    async def stop(self):
        """Stop the background flush task."""
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None
            # Flush remaining logs
            await self._flush_all()

    async def log(self, entry: LogEntry):
        """Queue a log entry."""
        await self._queue.put(entry)

    async def _flush_loop(self):
        """Background loop to flush logs periodically."""
        while True:
            try:
                await asyncio.sleep(LOG_FLUSH_INTERVAL)
                await self._flush_batch()
            except asyncio.CancelledError:
                raise

    async def _flush_batch(self):
        """Flush a batch of logs to database."""
        entries: List[LogEntry] = []

        # Collect up to batch_size entries
        while len(entries) < LOG_BATCH_SIZE:
            try:
                entry = self._queue.get_nowait()
                entries.append(entry)
            except asyncio.QueueEmpty:
                break

        if not entries:
            return

        # Insert batch
        async with db_connection() as db:
            await db.executemany(
                """INSERT INTO request_logs
                   (client_key_id, provider_id, model, prompt_tokens,
                    completion_tokens, latency_ms, success, error_message)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
                [
                    (
                        e.client_key_id, e.provider_id, e.model,
                        e.prompt_tokens, e.completion_tokens,
                        e.latency_ms, e.success, e.error_message
                    )
                    for e in entries
                ]
            )
            await db.commit()

    async def _flush_all(self):
        """Flush all remaining logs."""
        while not self._queue.empty():
            await self._flush_batch()


# Global instance
async_logger = AsyncLogger()