pinakes/crates/pinakes-core/src/jobs.rs

use std::{path::PathBuf, sync::Arc};

use chrono::{DateTime, Utc};
use rustc_hash::FxHashMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tokio::sync::{RwLock, mpsc};
use tokio_util::sync::CancellationToken;
use uuid::Uuid;

use crate::model::MediaId;

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", tag = "type")]
pub enum JobKind {
  Scan {
    path: Option<PathBuf>,
  },
  GenerateThumbnails {
    media_ids: Vec<MediaId>,
  },
  VerifyIntegrity {
    media_ids: Vec<MediaId>,
  },
  OrphanDetection,
  CleanupThumbnails,
  Export {
    format:      ExportFormat,
    destination: PathBuf,
  },
  Transcode {
    media_id: MediaId,
    profile:  String,
  },
  Enrich {
    media_ids: Vec<MediaId>,
  },
  CleanupAnalytics,
  TrashPurge,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ExportFormat {
  Json,
  Csv,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", tag = "state")]
pub enum JobStatus {
  Pending,
  Running { progress: f32, message: String },
  Completed { result: Value },
  Failed { error: String },
  Cancelled,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Job {
  pub id:         Uuid,
  pub kind:       JobKind,
  pub status:     JobStatus,
  pub created_at: DateTime<Utc>,
  pub updated_at: DateTime<Utc>,
}

struct WorkerItem {
  job_id: Uuid,
  kind:   JobKind,
  cancel: CancellationToken,
}

pub struct JobQueue {
  jobs:          Arc<RwLock<FxHashMap<Uuid, Job>>>,
  cancellations: Arc<RwLock<FxHashMap<Uuid, CancellationToken>>>,
  tx:            mpsc::Sender<WorkerItem>,
}

impl JobQueue {
  /// Create a new job queue and spawn `worker_count` background workers.
  ///
  /// The `executor` callback is invoked for each job; it receives the job kind,
  /// a progress-reporting callback, and a cancellation token.
  ///
  /// `job_timeout_secs` sets the maximum time a job can run before being
  /// cancelled. Set to 0 to disable the timeout.
  pub fn new<F>(
    worker_count: usize,
    job_timeout_secs: u64,
    executor: F,
  ) -> Arc<Self>
  where
    F: Fn(
        Uuid,
        JobKind,
        CancellationToken,
        Arc<RwLock<FxHashMap<Uuid, Job>>>,
      ) -> tokio::task::JoinHandle<()>
      + Send
      + Sync
      + 'static,
  {
    let (tx, rx) = mpsc::channel::<WorkerItem>(256);
    let rx = Arc::new(tokio::sync::Mutex::new(rx));
    let jobs: Arc<RwLock<FxHashMap<Uuid, Job>>> =
      Arc::new(RwLock::new(FxHashMap::default()));
    let cancellations: Arc<RwLock<FxHashMap<Uuid, CancellationToken>>> =
      Arc::new(RwLock::new(FxHashMap::default()));

    let executor = Arc::new(executor);

    for _ in 0..worker_count {
      let rx = Arc::clone(&rx);
      let jobs = Arc::clone(&jobs);
      let cancellations = Arc::clone(&cancellations);
      let executor = Arc::clone(&executor);

      tokio::spawn(async move {
        loop {
          let item = {
            let mut guard = rx.lock().await;
            guard.recv().await
          };
          let Some(item) = item else { break };

          // Mark as running
          {
            let mut map = jobs.write().await;
            if let Some(job) = map.get_mut(&item.job_id) {
              job.status = JobStatus::Running {
                progress: 0.0,
                message:  "starting".to_string(),
              };
              job.updated_at = Utc::now();
            }
          }

          let cancel_token = item.cancel.clone();
          let handle =
            executor(item.job_id, item.kind, item.cancel, Arc::clone(&jobs));

          if job_timeout_secs > 0 {
            let timeout = std::time::Duration::from_secs(job_timeout_secs);
            if tokio::time::timeout(timeout, handle).await.is_err() {
              // Timeout: cancel the job and mark as failed
              cancel_token.cancel();
              let mut map = jobs.write().await;
              if let Some(job) = map.get_mut(&item.job_id) {
                job.status = JobStatus::Failed {
                  error: format!("job timed out after {job_timeout_secs}s"),
                };
                job.updated_at = Utc::now();
              }
            }
          } else {
            let _ = handle.await;
          }

          // Clean up cancellation token
          cancellations.write().await.remove(&item.job_id);
        }
      });
    }

    Arc::new(Self {
      jobs,
      cancellations,
      tx,
    })
  }

  /// Submit a new job, returning its ID.
  pub async fn submit(&self, kind: JobKind) -> Uuid {
    let id = Uuid::now_v7();
    let now = Utc::now();
    let cancel = CancellationToken::new();

    let job = Job {
      id,
      kind: kind.clone(),
      status: JobStatus::Pending,
      created_at: now,
      updated_at: now,
    };

    {
      let mut map = self.jobs.write().await;
      map.insert(id, job);
      // Prune old terminal jobs to prevent unbounded memory growth.
      // Keep at most 500 completed/failed/cancelled entries, removing the
      // oldest.
      const MAX_TERMINAL_JOBS: usize = 500;
      let mut terminal: Vec<(Uuid, chrono::DateTime<Utc>)> = map
        .iter()
        .filter(|(_, j)| {
          matches!(
            j.status,
            JobStatus::Completed { .. }
              | JobStatus::Failed { .. }
              | JobStatus::Cancelled
          )
        })
        .map(|(k, j)| (*k, j.updated_at))
        .collect();
      if terminal.len() > MAX_TERMINAL_JOBS {
        terminal.sort_by_key(|(_, t)| *t);
        let to_remove = terminal.len() - MAX_TERMINAL_JOBS;
        for (stale_id, _) in terminal.into_iter().take(to_remove) {
          map.remove(&stale_id);
        }
      }
    }
    self.cancellations.write().await.insert(id, cancel.clone());

    let item = WorkerItem {
      job_id: id,
      kind,
      cancel,
    };

    // If the channel is full we still record the job; it will stay Pending
    let _ = self.tx.send(item).await;
    id
  }

  /// Get the status of a job.
  pub async fn status(&self, id: Uuid) -> Option<Job> {
    self.jobs.read().await.get(&id).cloned()
  }

  /// List all jobs, most recent first.
  pub async fn list(&self) -> Vec<Job> {
    let mut jobs: Vec<Job> = {
      let map = self.jobs.read().await;
      map.values().cloned().collect()
    };
    jobs.sort_by_key(|job| std::cmp::Reverse(job.created_at));
    jobs
  }

  /// Cancel a running or pending job.
  pub async fn cancel(&self, id: Uuid) -> bool {
    let token = {
      let guard = self.cancellations.read().await;
      guard.get(&id).cloned()
    };
    if let Some(token) = token {
      token.cancel();
      {
        let mut map = self.jobs.write().await;
        if let Some(job) = map.get_mut(&id) {
          job.status = JobStatus::Cancelled;
          job.updated_at = Utc::now();
        }
      }
      true
    } else {
      false
    }
  }

  /// Update a job's progress. Called by executors.
  pub async fn update_progress(
    jobs: &Arc<RwLock<FxHashMap<Uuid, Job>>>,
    id: Uuid,
    progress: f32,
    message: String,
  ) {
    let mut map = jobs.write().await;
    if let Some(job) = map.get_mut(&id) {
      job.status = JobStatus::Running { progress, message };
      job.updated_at = Utc::now();
    }
  }

  /// Mark a job as completed.
  pub async fn complete(
    jobs: &Arc<RwLock<FxHashMap<Uuid, Job>>>,
    id: Uuid,
    result: Value,
  ) {
    let mut map = jobs.write().await;
    if let Some(job) = map.get_mut(&id) {
      job.status = JobStatus::Completed { result };
      job.updated_at = Utc::now();
    }
  }

  /// Mark a job as failed.
  pub async fn fail(
    jobs: &Arc<RwLock<FxHashMap<Uuid, Job>>>,
    id: Uuid,
    error: String,
  ) {
    let mut map = jobs.write().await;
    if let Some(job) = map.get_mut(&id) {
      job.status = JobStatus::Failed { error };
      job.updated_at = Utc::now();
    }
  }

  /// Get job queue statistics
  pub async fn stats(&self) -> JobQueueStats {
    let jobs = self.jobs.read().await;
    let mut pending = 0;
    let mut running = 0;
    let mut completed = 0;
    let mut failed = 0;

    for job in jobs.values() {
      match job.status {
        JobStatus::Pending => pending += 1,
        JobStatus::Running { .. } => running += 1,
        JobStatus::Completed { .. } => completed += 1,
        JobStatus::Failed { .. } => failed += 1,
        JobStatus::Cancelled => {}, // Don't count cancelled jobs
      }
    }

    JobQueueStats {
      pending,
      running,
      completed,
      failed,
      total: jobs.len(),
    }
  }
}

/// Statistics about the job queue
#[derive(Debug, Clone, Default)]
pub struct JobQueueStats {
  pub pending:   usize,
  pub running:   usize,
  pub completed: usize,
  pub failed:    usize,
  pub total:     usize,
}