"""In-memory progress tracking for long-running background jobs.

This module is designed to support Server-Sent Events (SSE) streaming by keeping the
latest progress event for each job in memory and allowing subscribers to wait for
updates via an async generator.

Notes:
    - This is per-process memory. If you run multiple Uvicorn workers, each worker
      has its own progress store.
    - Jobs are automatically purged after completion/failure using a TTL to avoid
      memory leaks.
"""

from __future__ import annotations

import asyncio
import logging
import time
from collections.abc import AsyncGenerator
from typing import Any

from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)

_DEFAULT_TTL_SECONDS = 10 * 60


def _now_ms() -> int:
    """Return current time in milliseconds."""

    return int(time.time() * 1000)


class ProgressEvent(BaseModel):
    """Serializable progress payload.

    Attributes:
        job_id: Job identifier.
        stage: Current stage label (e.g., INITIALIZING, AI_EXTRACTION).
        progress: Percent (0-100).
        message: Human-readable status message.
        result: Optional final result payload (only when done).
        error: Optional error message (only when failed).
        done: Whether the job has reached a terminal state.
        timestamp_ms: Event creation timestamp.
    """

    job_id: str
    stage: str
    progress: int = Field(ge=0, le=100)
    message: str | None = None
    result: dict[str, Any] | None = None
    error: str | None = None
    done: bool = False
    timestamp_ms: int = Field(default_factory=_now_ms)


class _JobState:
    """Internal job state holder."""

    def __init__(self, event: ProgressEvent, owner_user_id: str | None = None) -> None:
        self.event = event
        self.owner_user_id = owner_user_id
        self.version = 0
        self.updated_at = time.monotonic()
        self.expires_at: float | None = None
        self.condition = asyncio.Condition()


class ProgressManager:
    """Track and publish job progress updates.

    This manager stores only the latest event per job. Subscribers receive the
    current snapshot immediately and then wait for further updates.
    """

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

    async def init_job(
        self,
        job_id: str,
        stage: str = "INITIALIZING",
        owner_user_id: str | None = None,
    ) -> ProgressEvent:
        """Create a new job if missing and return its current event.

        Args:
            job_id: Unique job id.
            stage: Initial stage label.
            owner_user_id: Optional owner user id for access control.

        Returns:
            ProgressEvent: The current (initial or existing) event.
        """

        async with self._lock:
            self._purge_expired_locked()
            existing = self._jobs.get(job_id)
            if existing is not None:
                if existing.owner_user_id is None and owner_user_id is not None:
                    existing.owner_user_id = owner_user_id
                return existing.event

            event = ProgressEvent(job_id=job_id, stage=stage, progress=0, message="Starting")
            self._jobs[job_id] = _JobState(event, owner_user_id=owner_user_id)
            logger.info(
                "Progress job initialized",
                extra={"job_id": job_id, "stage": stage, "has_owner": owner_user_id is not None},
            )
            return event

    async def snapshot(self, job_id: str) -> ProgressEvent:
        """Return the latest event snapshot for a job.

        Args:
            job_id: Job id.

        Returns:
            ProgressEvent: Latest event.

        Raises:
            KeyError: When job does not exist (or was purged).
        """

        async with self._lock:
            self._purge_expired_locked()
            state = self._jobs.get(job_id)
            if state is None:
                raise KeyError(job_id)
            return state.event

    async def update(
        self,
        job_id: str,
        stage: str,
        progress: int,
        message: str | None = None,
        result: dict[str, Any] | None = None,
        error: str | None = None,
        done: bool = False,
    ) -> ProgressEvent:
        """Update progress for a job and notify subscribers.

        Args:
            job_id: Job id.
            stage: Stage label.
            progress: Percent (0-100).
            message: Optional message.
            result: Optional result payload.
            error: Optional error message.
            done: Whether this update is terminal.

        Returns:
            ProgressEvent: The updated event.

        Raises:
            KeyError: When job is missing.
        """

        async with self._lock:
            self._purge_expired_locked()
            state = self._jobs.get(job_id)
            if state is None:
                raise KeyError(job_id)

        event = ProgressEvent(
            job_id=job_id,
            stage=stage,
            progress=progress,
            message=message,
            result=result,
            error=error,
            done=done,
        )

        await self._apply_event(job_id, event)
        return event

    async def complete(self, job_id: str, result: dict[str, Any] | None = None) -> ProgressEvent:
        """Mark a job as completed."""

        return await self.update(
            job_id=job_id,
            stage="COMPLETED",
            progress=100,
            message="Completed",
            result=result,
            done=True,
        )

    async def fail(self, job_id: str, error: str) -> ProgressEvent:
        """Mark a job as failed."""

        return await self.update(
            job_id=job_id,
            stage="FAILED",
            progress=100,
            message="Failed",
            error=error,
            done=True,
        )

    async def assert_owner(self, job_id: str, user_id: str) -> None:
        """Ensure the given user is the owner of the job.

        Args:
            job_id: Job id.
            user_id: Authenticated user id.

        Raises:
            KeyError: When job does not exist.
            PermissionError: When user is not the owner.
        """

        async with self._lock:
            self._purge_expired_locked()
            state = self._jobs.get(job_id)
            if state is None:
                raise KeyError(job_id)
            if state.owner_user_id is not None and state.owner_user_id != user_id:
                raise PermissionError(job_id)

    async def subscribe(self, job_id: str) -> AsyncGenerator[dict[str, Any], None]:
        """Subscribe to job progress updates.

        This yields the current snapshot immediately, then yields subsequent updates
        until a terminal event is received.

        Args:
            job_id: Job id.

        Yields:
            dict[str, Any]: JSON-serializable progress payload.

        Raises:
            KeyError: When job does not exist.
        """

        async with self._lock:
            self._purge_expired_locked()
            job = self._jobs.get(job_id)
            if job is None:
                raise KeyError(job_id)
            current_version = job.version
            event = job.event

        yield event.model_dump()

        while True:
            async with job.condition:
                while job.version == current_version:
                    await job.condition.wait()
                current_version = job.version
                event = job.event

            yield event.model_dump()

            if event.done:
                return

    async def _apply_event(self, job_id: str, event: ProgressEvent) -> None:
        """Apply an event to internal state and notify subscribers."""

        async with self._lock:
            self._purge_expired_locked()
            state = self._jobs.get(job_id)
            if state is None:
                raise KeyError(job_id)

            state.event = event
            state.version += 1
            state.updated_at = time.monotonic()
            if event.done:
                state.expires_at = state.updated_at + float(self._ttl_seconds)

        async with state.condition:
            state.condition.notify_all()

    def _purge_expired_locked(self) -> None:
        """Remove expired jobs.

        Must be called while holding `_lock`.
        """

        now = time.monotonic()
        expired = [
            job_id
            for job_id, state in self._jobs.items()
            if state.expires_at is not None and state.expires_at <= now
        ]
        for job_id in expired:
            self._jobs.pop(job_id, None)
            logger.info("Progress job purged", extra={"job_id": job_id})


# Module-level singleton used by routes/services.
progress_manager = ProgressManager()
