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

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

use chrono::{DateTime, Utc};
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<HashMap<Uuid, Job>>>,
  cancellations: Arc<RwLock<HashMap<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.
  pub fn new<F>(worker_count: usize, executor: F) -> Arc<Self>
  where
    F: Fn(
        Uuid,
        JobKind,
        CancellationToken,
        Arc<RwLock<HashMap<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<HashMap<Uuid, Job>>> =
      Arc::new(RwLock::new(HashMap::new()));
    let cancellations: Arc<RwLock<HashMap<Uuid, CancellationToken>>> =
      Arc::new(RwLock::new(HashMap::new()));

    let executor = Arc::new(executor);

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

      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 handle =
            executor(item.job_id, item.kind, item.cancel, jobs.clone());
          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,
    };

    self.jobs.write().await.insert(id, job);
    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 map = self.jobs.read().await;
    let mut jobs: Vec<Job> = 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 {
    if let Some(token) = self.cancellations.read().await.get(&id) {
      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<HashMap<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<HashMap<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<HashMap<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,
}