"""In-memory admission queue for the single-runner query gate.

Implements the state machine and estimate from docs/QUEUE_SPECIFICATION.md:
one runner slot that is idle, reserved for a single session, or running for
a single session, with a FIFO queue of waiting tickets behind it. All public
methods are synchronous and perform no suspension, so a single call is
atomic under asyncio's cooperative scheduling -- the same reasoning
backend.runner.QueryRunner relies on for its own gate.
"""

from __future__ import annotations

import secrets
import time
from dataclasses import dataclass
from typing import Callable, Dict, List, Optional

from backend.errors import QueueFullError, TicketNotFoundError

# Estimate smoothing factor and history window (spec section 4/5 constants).
ESTIMATE_ALPHA = 0.3
ESTIMATE_HISTORY = 20

# Runs cancelled by the client before this many seconds do not update the estimate.
CANCELLED_RUN_FLOOR_SECONDS = 5.0

TicketState = str  # "queued" | "ready" | "running"


@dataclass
class Ticket:
    ticket_id: str
    session_id: str
    joined_at: float
    last_seen_at: float
    state: TicketState
    reserved_at: Optional[float] = None
    expires_at: Optional[float] = None


def _generate_ticket_id() -> str:
    # Opaque, case-sensitive, [A-Za-z0-9_-] per CONTRACTS section 1.
    return "q_" + secrets.token_urlsafe(12)


class AdmissionController:
    """FIFO admission queue in front of the single query runner slot."""

    def __init__(
        self,
        reservation_seconds: float,
        heartbeat_seconds: float,
        max_queue_length: int,
        clock: Callable[[], float] = time.monotonic,
    ) -> None:
        self.reservation_seconds = reservation_seconds
        self.heartbeat_seconds = heartbeat_seconds
        self.max_queue_length = max_queue_length
        self.clock = clock

        self._tickets: Dict[str, Ticket] = {}
        self._queue: List[str] = []  # session_ids in "queued" state, FIFO
        self._reserved_session: Optional[str] = None
        self._running_session: Optional[str] = None
        self._running_started_at: Optional[float] = None

        self._run_durations: List[float] = []
        self._mean_duration: Optional[float] = None

    # -- Public API -----------------------------------------------------

    def join(self, session_id: str) -> dict:
        """Idempotent join. Returns the (possibly pre-existing) ticket view."""
        self._expire()
        existing = self._tickets.get(session_id)
        if existing is not None:
            return self._view(existing)

        now = self.clock()
        if self._reserved_session is None and self._running_session is None:
            # Runner is idle: book the place directly.
            ticket = Ticket(
                ticket_id=_generate_ticket_id(),
                session_id=session_id,
                joined_at=now,
                last_seen_at=now,
                state="ready",
                reserved_at=now,
                expires_at=now + self.reservation_seconds,
            )
            self._tickets[session_id] = ticket
            self._reserved_session = session_id
            return self._view(ticket)

        if len(self._queue) >= self.max_queue_length:
            raise QueueFullError("The queue is full; please try again later")

        ticket = Ticket(
            ticket_id=_generate_ticket_id(),
            session_id=session_id,
            joined_at=now,
            last_seen_at=now,
            state="queued",
        )
        self._tickets[session_id] = ticket
        self._queue.append(session_id)
        return self._view(ticket)

    def status(self, session_id: str) -> dict:
        """Heartbeat poll. Refreshes last_seen_at; raises TicketNotFoundError if absent."""
        self._expire()
        ticket = self._tickets.get(session_id)
        if ticket is None:
            raise TicketNotFoundError("No active queue ticket for this session")
        ticket.last_seen_at = self.clock()
        return self._view(ticket)

    def leave(self, session_id: str) -> None:
        """Drop the session's ticket, if any. Promotes the next ticket if it held the reservation."""
        self._expire()
        ticket = self._tickets.pop(session_id, None)
        if ticket is None:
            return
        if session_id in self._queue:
            self._queue.remove(session_id)
        if self._reserved_session == session_id:
            self._reserved_session = None
            self._promote()

    def claim(self, session_id: str) -> bool:
        """Attempt to move a query request's session into the running slot.

        Returns True if the session may proceed (claimed or implicitly joined),
        False if it must be refused with 409 busy. Does not disturb the queue
        on refusal.
        """
        now = self.clock()

        # The reserved session's own claim must see 409 once its window has
        # passed, even when this very call is what discovers the expiry --
        # it must not fall through to idle/implicit-join for itself.
        if self._reserved_session == session_id:
            ticket = self._tickets.get(session_id)
            expired = ticket is None or (
                ticket.expires_at is not None and now >= ticket.expires_at
            )
            if expired:
                self._tickets.pop(session_id, None)
                self._reserved_session = None
                self._promote()
                return False
            ticket.state = "running"
            ticket.reserved_at = None
            ticket.expires_at = None
            self._reserved_session = None
            self._running_session = session_id
            self._running_started_at = now
            return True

        # Not the current holder: evaluate lazy expiries (a stale reservation
        # held by someone else, or stale queued heartbeats) before deciding.
        self._expire()

        if self._running_session is not None:
            return False
        if self._reserved_session is not None:
            return False

        # Idle: implicit join, reserve, and claim in one step.
        ticket = Ticket(
            ticket_id=_generate_ticket_id(),
            session_id=session_id,
            joined_at=now,
            last_seen_at=now,
            state="running",
        )
        self._tickets[session_id] = ticket
        self._running_session = session_id
        self._running_started_at = now
        return True

    def run_end(self, session_id: str, duration_seconds: float, cancelled: bool = False) -> None:
        """Record run completion (success, failure, timeout, cancellation, disconnect).

        Must be called after the runner's own cleanup has finished so the next
        holder's reservation window never overlaps residual cleanup.
        """
        self._tickets.pop(session_id, None)
        if self._running_session == session_id:
            self._running_session = None
            self._running_started_at = None

        if not (cancelled and duration_seconds < CANCELLED_RUN_FLOOR_SECONDS):
            self._record_duration(duration_seconds)

        self._promote()

    def expire(self) -> None:
        """Public hook for periodic maintenance; evaluates lazy expiries eagerly."""
        self._expire()

    # -- Internal ---------------------------------------------------------

    def _record_duration(self, duration_seconds: float) -> None:
        self._run_durations.append(max(0.0, duration_seconds))
        if len(self._run_durations) > ESTIMATE_HISTORY:
            self._run_durations = self._run_durations[-ESTIMATE_HISTORY:]
        mean = self._run_durations[0]
        for d in self._run_durations[1:]:
            mean = ESTIMATE_ALPHA * d + (1 - ESTIMATE_ALPHA) * mean
        self._mean_duration = mean

    def _promote(self) -> None:
        if self._reserved_session is not None or self._running_session is not None:
            return
        if not self._queue:
            return
        session_id = self._queue.pop(0)
        ticket = self._tickets.get(session_id)
        if ticket is None:
            # Should not happen; keep state consistent and try the next one.
            self._promote()
            return
        now = self.clock()
        ticket.state = "ready"
        ticket.reserved_at = now
        ticket.expires_at = now + self.reservation_seconds
        self._reserved_session = session_id

    def _expire(self) -> None:
        now = self.clock()

        # Reservation expiry.
        if self._reserved_session is not None:
            ticket = self._tickets.get(self._reserved_session)
            if ticket is None or (ticket.expires_at is not None and now >= ticket.expires_at):
                if ticket is not None:
                    self._tickets.pop(self._reserved_session, None)
                self._reserved_session = None
                self._promote()

        # Heartbeat expiry for queued tickets only.
        stale = [
            sid
            for sid in self._queue
            if (ticket := self._tickets.get(sid)) is not None
            and now - ticket.last_seen_at > self.heartbeat_seconds
        ]
        for sid in stale:
            self._tickets.pop(sid, None)
            if sid in self._queue:
                self._queue.remove(sid)

    def _position(self, session_id: str, ticket: Ticket) -> int:
        if ticket.state == "running":
            return 0
        front: List[str] = []
        if self._reserved_session is not None:
            front.append(self._reserved_session)
        front.extend(self._queue)
        try:
            return front.index(session_id) + 1
        except ValueError:
            return 0

    def _eta_seconds(self, ahead: int) -> Optional[float]:
        if self._mean_duration is None:
            return None
        if self._running_session is not None:
            started_at = self._running_started_at if self._running_started_at is not None else self.clock()
            elapsed = self.clock() - started_at
            return max(0.0, self._mean_duration * (ahead + 1) - elapsed)
        return max(0.0, self._mean_duration * (ahead + 1))

    def _view(self, ticket: Ticket) -> dict:
        position = self._position(ticket.session_id, ticket)
        ahead = max(0, position - 1)
        eta = self._eta_seconds(ahead)
        runner = "running" if self._running_session is not None else "reserved"

        reservation_expires_in: Optional[float] = None
        if ticket.state == "ready" and ticket.expires_at is not None:
            reservation_expires_in = max(0.0, ticket.expires_at - self.clock())

        return {
            "ticket_id": ticket.ticket_id,
            "status": ticket.state,
            "position": position,
            "ahead": ahead,
            "eta_seconds": round(eta) if eta is not None else None,
            "reservation_expires_in_seconds": (
                round(reservation_expires_in) if reservation_expires_in is not None else None
            ),
            "runner": runner,
        }