"""Shared in-memory job store for background tasks (training, optimisation).

Fine for a single-server dev tool.  Jobs older than ``ttl_seconds`` are
evicted on each ``create_job`` / ``get_job`` call to bound memory usage.
"""

from __future__ import annotations

import time
import uuid
from typing import Any

from fastapi import HTTPException

_DEFAULT_TTL_SECONDS = 24 * 60 * 60  # 24 hours


class JobStore:
    """Thread-safe-enough dict-backed job store with TTL eviction.

    Each route module creates its own instance so job-ID namespaces stay
    independent (a training job ID will never collide with an optimiser
    job ID, just as before the refactor).
    """

    def __init__(self, ttl_seconds: int = _DEFAULT_TTL_SECONDS) -> None:
        self._jobs: dict[str, dict[str, Any]] = {}
        self._ttl_seconds = ttl_seconds

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _evict_stale(self) -> None:
        """Remove jobs older than TTL to bound memory usage."""
        cutoff = time.time() - self._ttl_seconds
        stale = [
            jid
            for jid, j in self._jobs.items()
            if j.get("created_at", 0) < cutoff and j.get("status") not in ("running",)
        ]
        for jid in stale:
            del self._jobs[jid]

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

    def create_job(self, initial_status: dict[str, Any]) -> str:
        """Generate a UUID, store *initial_status* with a timestamp, return the ID.

        Automatically evicts stale jobs before inserting.
        """
        self._evict_stale()
        job_id = uuid.uuid4().hex[:12]
        initial_status.setdefault("created_at", time.time())
        self._jobs[job_id] = initial_status
        return job_id

    def get_job(self, job_id: str) -> dict[str, Any] | None:
        """Return the job dict for *job_id*, or ``None`` if not found.

        Evicts stale jobs first so callers never see expired entries.
        """
        self._evict_stale()
        return self._jobs.get(job_id)

    def update_job(self, job_id: str, **fields: Any) -> None:
        """Merge *fields* into the stored job dict.

        Raises ``KeyError`` if *job_id* does not exist.
        """
        self._jobs[job_id].update(fields)

    def atomic_update(
        self,
        job_id: str,
        fields: dict[str, Any],
        *,
        expected_status: str | None = None,
    ) -> dict[str, Any]:
        """Replace the job dict with a merged copy — thread-safe.

        Instead of mutating the existing dict (which can race with
        concurrent readers), this builds a **new** dict and swaps it in
        with a single pointer assignment.  CPython's GIL guarantees that
        ``dict.__setitem__`` is atomic, so a reader will always see
        either the old dict or the new one — never a half-updated state.

        When *expected_status* is provided, the update is skipped if the
        current status does not match (prevents timeout from overwriting
        a completed job).

        Raises ``KeyError`` if *job_id* does not exist.
        """
        old = self._jobs[job_id]
        if expected_status is not None and old.get("status") != expected_status:
            return old
        self._jobs[job_id] = {**old, **fields}
        return self._jobs[job_id]

    def require_job(self, job_id: str) -> dict[str, Any]:
        """Return the job dict for *job_id*, or raise HTTP 404 if not found.

        Convenience wrapper around :meth:`get_job` that eliminates the
        repetitive ``if job is None: raise HTTPException(...)`` guard at
        every call site.
        """
        job = self.get_job(job_id)
        if job is None:
            raise HTTPException(status_code=404, detail=f"Job '{job_id}' not found")
        return job

    def require_completed_job(self, job_id: str) -> dict[str, Any]:
        """Return the job dict for *job_id*, raising if missing or not completed.

        Combines :meth:`require_job` (404 if not found) with a status
        check (400 if not ``"completed"``).  Eliminates the repetitive
        two-step guard pattern at call sites that need a finished job.
        """
        job = self.require_job(job_id)
        if job.get("status") != "completed":
            raise HTTPException(
                status_code=400,
                detail=f"Job '{job_id}' is not completed (status: {job.get('status')})",
            )
        return job

    def clear_result_data(
        self,
        job_id: str,
        keys: tuple[str, ...] = ("solver", "solve_result", "quote_grid"),
    ) -> None:
        """Remove heavy objects from a completed job to free memory.

        After a solve result has been saved or logged to MLflow, the full
        solver, solve result (entire scored DataFrame), and QuoteGrid are
        no longer needed.  This method strips those keys from the job dict
        while keeping lightweight metadata (status, config, result summary)
        intact, allowing the TTL eviction to work on a much smaller dict.

        No-op if *job_id* does not exist or keys are already absent.
        """
        job = self._jobs.get(job_id)
        if job is None:
            return
        cleaned = {k: v for k, v in job.items() if k not in keys}
        self._jobs[job_id] = cleaned

    @property
    def jobs(self) -> dict[str, dict[str, Any]]:
        """Direct access to the underlying dict.

        Provided for callsites that need to iterate (e.g. checking for
        running jobs).  Prefer ``get_job`` / ``update_job`` for single-key
        access.
        """
        return self._jobs
