jobs.py

"""
File-based job queue for long-running evaluation tasks.

This module provides persistent job storage that survives HF Spaces timeouts
and restarts. Jobs are stored as individual JSON files for atomic updates
and to avoid file locking issues with concurrent access.

Architecture:
- Each job is stored as a separate file: jobs/{job_id}.json
- Background threads write progress updates to these files
- The Gradio timer polls by reading these files (stateless, fast)
- Old jobs are cleaned up periodically
"""

import json
import os
import time
import uuid
from dataclasses import dataclass, asdict
from datetime import datetime, timedelta
from pathlib import Path
from typing import Optional, Any


# Use /data on HF Spaces (persistent storage), fallback to local directory
if os.path.exists("/data"):
    JOBS_DIR = Path("/data/jobs")
else:
    JOBS_DIR = Path("jobs")


@dataclass
class Job:
    """Represents an evaluation job."""

    job_id: str
    model_id: str
    model_name: str
    max_pairs: int
    status: str  # "pending", "running", "completed", "failed"
    progress: str  # Human-readable progress message
    progress_pct: float  # 0.0 to 1.0
    result: Optional[str]  # Markdown result when completed
    error: Optional[str]  # Error message if failed
    created_at: str  # ISO timestamp
    updated_at: str  # ISO timestamp

    def to_dict(self) -> dict:
        """Convert to dictionary for JSON serialization."""
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict) -> "Job":
        """Create from dictionary."""
        return cls(**data)


def _ensure_jobs_dir():
    """Ensure the jobs directory exists."""
    JOBS_DIR.mkdir(parents=True, exist_ok=True)


def _job_path(job_id: str) -> Path:
    """Get the file path for a job."""
    return JOBS_DIR / f"{job_id}.json"


def create_job(model_id: str, model_name: str, max_pairs: int) -> Job:
    """
    Create a new job and save it to disk.

    Args:
        model_id: The model identifier
        model_name: Display name for the model
        max_pairs: Maximum number of pairs to evaluate

    Returns:
        The created Job object
    """
    _ensure_jobs_dir()

    now = datetime.now().isoformat()
    job = Job(
        job_id=str(uuid.uuid4()),
        model_id=model_id,
        model_name=model_name,
        max_pairs=max_pairs,
        status="pending",
        progress="⏳ Job created, waiting to start...",
        progress_pct=0.0,
        result=None,
        error=None,
        created_at=now,
        updated_at=now,
    )

    _save_job(job)
    return job


def _save_job(job: Job):
    """
    Save a job to disk atomically.

    Uses write-to-temp-then-rename pattern to prevent corruption
    from concurrent reads or crashes during write.
    """
    _ensure_jobs_dir()

    job_path = _job_path(job.job_id)
    temp_path = job_path.with_suffix(".tmp")

    # Write to temp file first
    temp_path.write_text(json.dumps(job.to_dict(), indent=2))

    # Atomic rename (on POSIX systems)
    temp_path.rename(job_path)


def get_job(job_id: str) -> Optional[Job]:
    """
    Get a job by ID.

    Args:
        job_id: The job identifier

    Returns:
        Job object if found, None otherwise
    """
    job_path = _job_path(job_id)

    if not job_path.exists():
        return None

    try:
        data = json.loads(job_path.read_text())
        return Job.from_dict(data)
    except (json.JSONDecodeError, TypeError, KeyError) as e:
        # Corrupted job file - remove it
        print(f"Warning: Corrupted job file {job_id}, removing: {e}")
        try:
            job_path.unlink()
        except OSError:
            pass
        return None


def update_job_progress(
    job_id: str,
    progress: str,
    progress_pct: float,
    status: str = "running",
):
    """
    Update job progress. Called frequently by background thread.

    Args:
        job_id: The job identifier
        progress: Human-readable progress message
        progress_pct: Progress percentage (0.0 to 1.0)
        status: Job status (usually "running")
    """
    job = get_job(job_id)
    if job is None:
        return

    job.status = status
    job.progress = progress
    job.progress_pct = progress_pct
    job.updated_at = datetime.now().isoformat()

    _save_job(job)


def complete_job(job_id: str, result: str):
    """
    Mark a job as completed with results.

    Args:
        job_id: The job identifier
        result: Markdown-formatted result string
    """
    job = get_job(job_id)
    if job is None:
        return

    job.status = "completed"
    job.progress = "✅ Evaluation complete!"
    job.progress_pct = 1.0
    job.result = result
    job.updated_at = datetime.now().isoformat()

    _save_job(job)


def fail_job(job_id: str, error: str):
    """
    Mark a job as failed with an error message.

    Args:
        job_id: The job identifier
        error: Error message
    """
    job = get_job(job_id)
    if job is None:
        return

    job.status = "failed"
    job.progress = f"❌ Error: {error}"
    job.progress_pct = 0.0
    job.error = error
    job.updated_at = datetime.now().isoformat()

    _save_job(job)


def delete_job(job_id: str):
    """
    Delete a job file.

    Args:
        job_id: The job identifier
    """
    job_path = _job_path(job_id)
    try:
        job_path.unlink(missing_ok=True)
    except OSError:
        pass


def cleanup_old_jobs(max_age_hours: int = 24):
    """
    Clean up jobs older than the specified age.

    Args:
        max_age_hours: Maximum age in hours before cleanup
    """
    _ensure_jobs_dir()

    cutoff = datetime.now() - timedelta(hours=max_age_hours)

    for job_path in JOBS_DIR.glob("*.json"):
        try:
            data = json.loads(job_path.read_text())
            updated_at = datetime.fromisoformat(data.get("updated_at", ""))

            if updated_at < cutoff:
                print(f"Cleaning up old job: {job_path.name}")
                job_path.unlink()
        except (json.JSONDecodeError, ValueError, OSError) as e:
            # If we can't parse it, it's probably corrupted - remove it
            print(f"Removing unparseable job file {job_path.name}: {e}")
            try:
                job_path.unlink()
            except OSError:
                pass


def cleanup_stale_jobs(stale_minutes: int = 30):
    """
    Clean up jobs that appear to be stale (stuck in running state).

    This handles cases where the Space restarted mid-evaluation.

    Args:
        stale_minutes: Minutes without update before considering stale
    """
    _ensure_jobs_dir()

    cutoff = datetime.now() - timedelta(minutes=stale_minutes)

    for job_path in JOBS_DIR.glob("*.json"):
        try:
            data = json.loads(job_path.read_text())
            status = data.get("status", "")
            updated_at = datetime.fromisoformat(data.get("updated_at", ""))

            # Only clean up jobs that are stuck in running/pending state
            if status in ("running", "pending") and updated_at < cutoff:
                print(f"Marking stale job as failed: {job_path.name}")
                job_id = data.get("job_id")
                if job_id:
                    fail_job(job_id, "Job timed out or server restarted. Please try again.")
        except (json.JSONDecodeError, ValueError, OSError) as e:
            pass


def list_jobs(status: Optional[str] = None, limit: int = 100) -> list[Job]:
    """
    List jobs, optionally filtered by status.

    Args:
        status: Filter by status (None for all)
        limit: Maximum number of jobs to return

    Returns:
        List of Job objects, sorted by created_at descending
    """
    _ensure_jobs_dir()

    jobs = []
    for job_path in JOBS_DIR.glob("*.json"):
        try:
            data = json.loads(job_path.read_text())
            job = Job.from_dict(data)

            if status is None or job.status == status:
                jobs.append(job)
        except (json.JSONDecodeError, TypeError, KeyError):
            pass

    # Sort by created_at descending
    jobs.sort(key=lambda j: j.created_at, reverse=True)

    return jobs[:limit]


def get_active_job_count() -> int:
    """
    Get the count of currently active (running/pending) jobs.

    Returns:
        Number of active jobs
    """
    return len(list_jobs(status="running")) + len(list_jobs(status="pending"))