//! Enhance jobs (#192): one row per ComfyUI run submitted for a photo.
//!
//! The row is the only state: the poller advances it from ComfyUI's
//! `/queue` and `/history`, the endpoints read it, approve appends to
//! `saved`. `pending` / `saved` hold small JSON documents owned by
//! `crate::enhance`; this layer treats them as opaque text.

use std::sync::{Arc, Mutex};

use diesel::prelude::*;
use diesel::sql_query;
use diesel::sql_types::{BigInt, Integer, Nullable, Text};
use diesel::sqlite::SqliteConnection;

use crate::database::{DbError, DbErrorKind, connect};
use crate::otel::trace_db_call;

pub const STATUS_QUEUED: &str = "queued";
pub const STATUS_RUNNING: &str = "running";
pub const STATUS_READY: &str = "ready";
pub const STATUS_APPROVED: &str = "approved";
pub const STATUS_FAILED: &str = "failed";
pub const STATUS_EXPIRED: &str = "expired";
pub const STATUS_UNAVAILABLE: &str = "unavailable";

#[derive(Debug, Clone, PartialEq, Eq, QueryableByName)]
pub struct EnhanceJobRow {
    #[diesel(sql_type = Text)]
    pub id: String,
    #[diesel(sql_type = Text)]
    pub original_hash: String,
    #[diesel(sql_type = Integer)]
    pub library_id: i32,
    #[diesel(sql_type = Text)]
    pub rel_path: String,
    #[diesel(sql_type = Text)]
    pub template: String,
    #[diesel(sql_type = Nullable<Text>)]
    pub prompt: Option<String>,
    #[diesel(sql_type = Nullable<Text>)]
    pub reference_hash: Option<String>,
    #[diesel(sql_type = Nullable<Text>)]
    pub reference_path: Option<String>,
    #[diesel(sql_type = Text)]
    pub comfy_prompt_id: String,
    #[diesel(sql_type = Text)]
    pub status: String,
    #[diesel(sql_type = Nullable<Text>)]
    pub error: Option<String>,
    #[diesel(sql_type = Nullable<Integer>)]
    pub queue_position: Option<i32>,
    #[diesel(sql_type = Text)]
    pub pending: String,
    #[diesel(sql_type = Text)]
    pub saved: String,
    #[diesel(sql_type = Integer)]
    pub discarded: i32,
    #[diesel(sql_type = BigInt)]
    pub created_at: i64,
    #[diesel(sql_type = BigInt)]
    pub updated_at: i64,
}

const COLUMNS: &str = "id, original_hash, library_id, rel_path, template, prompt, reference_hash, \
     reference_path, comfy_prompt_id, status, error, queue_position, pending, saved, \
     discarded, created_at, updated_at";

/// What the poller / approve changes on an existing job.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct JobUpdate {
    pub status: String,
    pub error: Option<String>,
    pub queue_position: Option<i32>,
    /// `Some` replaces the `pending` document.
    pub pending: Option<String>,
    /// `Some` replaces the `saved` document.
    pub saved: Option<String>,
}

impl JobUpdate {
    pub fn status(status: &str) -> Self {
        JobUpdate {
            status: status.to_string(),
            error: None,
            queue_position: None,
            pending: None,
            saved: None,
        }
    }
}

pub trait EnhanceJobDao: Sync + Send {
    fn insert(&mut self, job: &EnhanceJobRow) -> Result<(), DbError>;
    fn get(&mut self, id: &str) -> Result<Option<EnhanceJobRow>, DbError>;
    /// Newest first.
    fn list_for_hash(
        &mut self,
        original_hash: &str,
        limit: i64,
    ) -> Result<Vec<EnhanceJobRow>, DbError>;
    /// Jobs the poller still has to advance (`queued` / `running`).
    fn list_active(&mut self) -> Result<Vec<EnhanceJobRow>, DbError>;
    fn update(&mut self, id: &str, update: &JobUpdate) -> Result<(), DbError>;
    /// Hide a job from `list_for_hash`.
    fn discard(&mut self, id: &str) -> Result<(), DbError>;
    /// Delete jobs last touched before `cutoff` (unix seconds) that are not
    /// active. Returns the number removed.
    fn prune(&mut self, cutoff: i64) -> Result<usize, DbError>;
}

pub struct SqliteEnhanceJobDao {
    connection: Arc<Mutex<SqliteConnection>>,
}

impl Default for SqliteEnhanceJobDao {
    fn default() -> Self {
        Self::new()
    }
}

impl SqliteEnhanceJobDao {
    pub fn new() -> Self {
        SqliteEnhanceJobDao {
            connection: Arc::new(Mutex::new(connect())),
        }
    }

    /// Test-only: share a connection with other DAOs.
    #[cfg(test)]
    pub fn from_shared(connection: Arc<Mutex<SqliteConnection>>) -> Self {
        SqliteEnhanceJobDao { connection }
    }
}

impl EnhanceJobDao for SqliteEnhanceJobDao {
    fn insert(&mut self, job: &EnhanceJobRow) -> Result<(), DbError> {
        trace_db_call("insert", "enhance_job_insert", |_span| {
            let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
            sql_query(format!(
                "INSERT INTO enhance_jobs ({COLUMNS}) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)"
            ))
            .bind::<Text, _>(&job.id)
            .bind::<Text, _>(&job.original_hash)
            .bind::<Integer, _>(job.library_id)
            .bind::<Text, _>(&job.rel_path)
            .bind::<Text, _>(&job.template)
            .bind::<Nullable<Text>, _>(&job.prompt)
            .bind::<Nullable<Text>, _>(&job.reference_hash)
            .bind::<Nullable<Text>, _>(&job.reference_path)
            .bind::<Text, _>(&job.comfy_prompt_id)
            .bind::<Text, _>(&job.status)
            .bind::<Nullable<Text>, _>(&job.error)
            .bind::<Nullable<Integer>, _>(job.queue_position)
            .bind::<Text, _>(&job.pending)
            .bind::<Text, _>(&job.saved)
            .bind::<Integer, _>(job.discarded)
            .bind::<BigInt, _>(job.created_at)
            .bind::<BigInt, _>(job.updated_at)
            .execute(&mut *conn)
            .map(|_| ())
            .map_err(|e| anyhow::anyhow!("Insert error: {}", e))
        })
        .map_err(|e| DbError::log(DbErrorKind::InsertError, e))
    }

    fn get(&mut self, id: &str) -> Result<Option<EnhanceJobRow>, DbError> {
        trace_db_call("query", "enhance_job_get", |_span| {
            let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
            sql_query(format!("SELECT {COLUMNS} FROM enhance_jobs WHERE id = ?"))
                .bind::<Text, _>(id)
                .load::<EnhanceJobRow>(&mut *conn)
                .map(|rows| rows.into_iter().next())
                .map_err(|e| anyhow::anyhow!("Query error: {}", e))
        })
        .map_err(|e| DbError::log(DbErrorKind::QueryError, e))
    }

    fn list_for_hash(
        &mut self,
        original_hash: &str,
        limit: i64,
    ) -> Result<Vec<EnhanceJobRow>, DbError> {
        trace_db_call("query", "enhance_job_list_for_hash", |_span| {
            let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
            sql_query(format!(
                "SELECT {COLUMNS} FROM enhance_jobs WHERE original_hash = ? AND discarded = 0 \
                 ORDER BY created_at DESC, rowid DESC LIMIT ?"
            ))
            .bind::<Text, _>(original_hash)
            .bind::<BigInt, _>(limit)
            .load::<EnhanceJobRow>(&mut *conn)
            .map_err(|e| anyhow::anyhow!("Query error: {}", e))
        })
        .map_err(|e| DbError::log(DbErrorKind::QueryError, e))
    }

    fn list_active(&mut self) -> Result<Vec<EnhanceJobRow>, DbError> {
        trace_db_call("query", "enhance_job_list_active", |_span| {
            let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
            sql_query(format!(
                "SELECT {COLUMNS} FROM enhance_jobs WHERE status IN ('queued', 'running') \
                 ORDER BY created_at, rowid"
            ))
            .load::<EnhanceJobRow>(&mut *conn)
            .map_err(|e| anyhow::anyhow!("Query error: {}", e))
        })
        .map_err(|e| DbError::log(DbErrorKind::QueryError, e))
    }

    fn update(&mut self, id: &str, u: &JobUpdate) -> Result<(), DbError> {
        trace_db_call("update", "enhance_job_update", |_span| {
            let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
            sql_query(
                "UPDATE enhance_jobs SET status = ?, error = ?, queue_position = ?, \
                 pending = COALESCE(?, pending), saved = COALESCE(?, saved), updated_at = ? \
                 WHERE id = ?",
            )
            .bind::<Text, _>(&u.status)
            .bind::<Nullable<Text>, _>(&u.error)
            .bind::<Nullable<Integer>, _>(u.queue_position)
            .bind::<Nullable<Text>, _>(&u.pending)
            .bind::<Nullable<Text>, _>(&u.saved)
            .bind::<BigInt, _>(chrono::Utc::now().timestamp())
            .bind::<Text, _>(id)
            .execute(&mut *conn)
            .map(|_| ())
            .map_err(|e| anyhow::anyhow!("Update error: {}", e))
        })
        .map_err(|e| DbError::log(DbErrorKind::UpdateError, e))
    }

    fn discard(&mut self, id: &str) -> Result<(), DbError> {
        trace_db_call("update", "enhance_job_discard", |_span| {
            let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
            sql_query("UPDATE enhance_jobs SET discarded = 1, updated_at = ? WHERE id = ?")
                .bind::<BigInt, _>(chrono::Utc::now().timestamp())
                .bind::<Text, _>(id)
                .execute(&mut *conn)
                .map(|_| ())
                .map_err(|e| anyhow::anyhow!("Update error: {}", e))
        })
        .map_err(|e| DbError::log(DbErrorKind::UpdateError, e))
    }

    fn prune(&mut self, cutoff: i64) -> Result<usize, DbError> {
        trace_db_call("delete", "enhance_job_prune", |_span| {
            let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
            sql_query(
                "DELETE FROM enhance_jobs WHERE updated_at < ? \
                 AND status NOT IN ('queued', 'running')",
            )
            .bind::<BigInt, _>(cutoff)
            .execute(&mut *conn)
            .map_err(|e| anyhow::anyhow!("Delete error: {}", e))
        })
        .map_err(|e| DbError::log(DbErrorKind::UpdateError, e))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::database::test::in_memory_db_connection;

    fn dao() -> SqliteEnhanceJobDao {
        SqliteEnhanceJobDao::from_shared(Arc::new(Mutex::new(in_memory_db_connection())))
    }

    fn job(id: &str, hash: &str, status: &str, created: i64) -> EnhanceJobRow {
        EnhanceJobRow {
            id: id.into(),
            original_hash: hash.into(),
            library_id: 1,
            rel_path: "a.jpg".into(),
            template: "restore".into(),
            prompt: Some("warmer".into()),
            reference_hash: None,
            reference_path: None,
            comfy_prompt_id: format!("p-{id}"),
            status: status.into(),
            error: None,
            queue_position: None,
            pending: "[]".into(),
            saved: "[]".into(),
            discarded: 0,
            created_at: created,
            updated_at: created,
        }
    }

    #[test]
    fn insert_get_list_update_prune() {
        let mut d = dao();
        d.insert(&job("a", "h1", "queued", 100)).unwrap();
        d.insert(&job("b", "h1", "ready", 200)).unwrap();
        d.insert(&job("c", "h2", "running", 300)).unwrap();

        assert_eq!(d.get("a").unwrap().unwrap().comfy_prompt_id, "p-a");
        assert!(d.get("zzz").unwrap().is_none());

        let ids: Vec<String> = d
            .list_for_hash("h1", 10)
            .unwrap()
            .into_iter()
            .map(|j| j.id)
            .collect();
        assert_eq!(ids, ["b", "a"], "newest first");
        assert_eq!(d.list_for_hash("h1", 1).unwrap().len(), 1);

        let active: Vec<String> = d.list_active().unwrap().into_iter().map(|j| j.id).collect();
        assert_eq!(active, ["a", "c"]);

        d.update(
            "a",
            &JobUpdate {
                status: "ready".into(),
                error: None,
                queue_position: None,
                pending: Some("[{\"id\":\"x\"}]".into()),
                saved: None,
            },
        )
        .unwrap();
        let a = d.get("a").unwrap().unwrap();
        assert_eq!(a.status, "ready");
        assert_eq!(a.pending, "[{\"id\":\"x\"}]");
        assert_eq!(a.saved, "[]", "unset columns are left alone");
        assert!(a.updated_at >= 100);

        d.discard("b").unwrap();
        assert_eq!(
            d.list_for_hash("h1", 10).unwrap().len(),
            1,
            "discarded hidden"
        );

        // Prune drops only old, finished jobs (updated_at was just bumped
        // for `a`; `c` is active).
        assert_eq!(d.prune(i64::MAX).unwrap(), 2);
        assert!(d.get("c").unwrap().is_some());
        assert!(d.get("b").unwrap().is_none());
    }
}