use anyhow::Result;
use rusqlite::{params, Connection};
use serde::Serialize;

use crate::db::ExtractionTaskKind;

#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct ExtractionReplayRange {
    pub id: i64,
    pub source_task_id: i64,
    pub replay_task_id: Option<i64>,
    pub task_kind: String,
    pub project: String,
    pub session_id: Option<String>,
    pub from_event_id: i64,
    pub to_event_id: i64,
    pub status: String,
    pub attempts: i64,
    pub updated_at_epoch: i64,
    pub last_error: Option<String>,
}

pub fn list_extraction_replay_ranges(
    conn: &Connection,
    project: Option<&str>,
    limit: i64,
) -> Result<Vec<ExtractionReplayRange>> {
    let mut stmt = conn.prepare(
        "SELECT r.id, r.source_task_id, r.replay_task_id, r.task_kind, p.project_path,
                s.session_id, r.from_event_id, r.to_event_id, r.status, r.attempts,
                r.updated_at_epoch, r.last_error
         FROM extraction_replay_ranges r
         JOIN projects p ON p.id = r.project_id
         LEFT JOIN sessions s ON s.id = r.session_row_id
         WHERE r.status IN ('pending', 'failed', 'requeued', 'quarantined')
           AND (?1 IS NULL OR p.project_path = ?1)
         ORDER BY r.updated_at_epoch DESC, r.id DESC
         LIMIT ?2",
    )?;
    let rows = stmt.query_map(params![project, limit.max(1)], |row| {
        Ok(ExtractionReplayRange {
            id: row.get(0)?,
            source_task_id: row.get(1)?,
            replay_task_id: row.get(2)?,
            task_kind: row.get(3)?,
            project: row.get(4)?,
            session_id: row.get(5)?,
            from_event_id: row.get(6)?,
            to_event_id: row.get(7)?,
            status: row.get(8)?,
            attempts: row.get(9)?,
            updated_at_epoch: row.get(10)?,
            last_error: row.get(11)?,
        })
    })?;
    crate::db::query::collect_rows(rows)
}

pub fn count_retryable_extraction_replay_ranges(
    conn: &Connection,
    project: Option<&str>,
    limit: i64,
) -> Result<i64> {
    conn.query_row(
        "SELECT COUNT(*)
         FROM (
             SELECT r.id
             FROM extraction_replay_ranges r
             JOIN projects p ON p.id = r.project_id
             WHERE r.status IN ('pending', 'failed')
               AND NOT EXISTS (
                 SELECT 1
                 FROM extraction_tasks t
                 WHERE t.replay_range_id = r.id
                   AND t.status IN ('pending', 'processing')
               )
               AND (?1 IS NULL OR p.project_path = ?1)
             ORDER BY r.updated_at_epoch ASC, r.id ASC
             LIMIT ?2
         )",
        params![project, limit.max(1)],
        |row| row.get(0),
    )
    .map_err(Into::into)
}

fn query_retryable_replay_range_ids(
    conn: &Connection,
    project: Option<&str>,
    limit: i64,
) -> Result<Vec<i64>> {
    let mut stmt = conn.prepare(
        "SELECT r.id
         FROM extraction_replay_ranges r
         JOIN projects p ON p.id = r.project_id
         WHERE r.status IN ('pending', 'failed')
           AND NOT EXISTS (
             SELECT 1
             FROM extraction_tasks t
             WHERE t.replay_range_id = r.id
               AND t.status IN ('pending', 'processing')
           )
           AND (?1 IS NULL OR p.project_path = ?1)
         ORDER BY r.updated_at_epoch ASC, r.id ASC
         LIMIT ?2",
    )?;
    let rows = stmt.query_map(params![project, limit.max(1)], |row| row.get::<_, i64>(0))?;
    crate::db::query::collect_rows(rows)
}

pub fn retry_extraction_replay_ranges(
    conn: &Connection,
    project: Option<&str>,
    limit: i64,
) -> Result<usize> {
    let tx = conn.unchecked_transaction()?;
    let range_ids = query_retryable_replay_range_ids(&tx, project, limit)?;
    for range_id in &range_ids {
        enqueue_replay_extraction_task(&tx, *range_id)?;
    }
    tx.commit()?;
    Ok(range_ids.len())
}

pub fn quarantine_extraction_replay_ranges(
    conn: &Connection,
    project: Option<&str>,
    limit: i64,
) -> Result<usize> {
    let tx = conn.unchecked_transaction()?;
    let range_ids = query_retryable_replay_range_ids(&tx, project, limit)?;
    let now = chrono::Utc::now().timestamp();
    for range_id in &range_ids {
        tx.execute(
            "UPDATE extraction_replay_ranges
             SET status = 'quarantined', updated_at_epoch = ?1
             WHERE id = ?2",
            params![now, range_id],
        )?;
        clear_terminal_failures_for_quiesced_range(&tx, *range_id, now)?;
    }
    tx.commit()?;
    Ok(range_ids.len())
}

fn enqueue_replay_extraction_task(conn: &Connection, range_id: i64) -> Result<i64> {
    let (task_kind, host_id, workspace_id, project_id, session_row_id, from_event_id, to_event_id) =
        conn.query_row(
            "SELECT task_kind, host_id, workspace_id, project_id, session_row_id,
                    from_event_id, to_event_id
             FROM extraction_replay_ranges
             WHERE id = ?1 AND status IN ('pending', 'failed')",
            params![range_id],
            |row| {
                Ok((
                    row.get::<_, String>(0)?,
                    row.get::<_, i64>(1)?,
                    row.get::<_, i64>(2)?,
                    row.get::<_, i64>(3)?,
                    row.get::<_, Option<i64>>(4)?,
                    row.get::<_, i64>(5)?,
                    row.get::<_, i64>(6)?,
                ))
            },
        )?;
    let session_row_id = session_row_id.ok_or_else(|| {
        anyhow::anyhow!("extraction replay range {range_id} is missing session_row_id")
    })?;
    let task_kind_value = ExtractionTaskKind::from_db(&task_kind)?;
    let now = chrono::Utc::now().timestamp();
    let idempotency_key =
        format!("{host_id}:{project_id}:{session_row_id}:{task_kind}:replay-range:{range_id}");
    conn.execute(
        "INSERT INTO extraction_tasks
         (task_kind, host_id, workspace_id, project_id, session_row_id, priority, status,
          idempotency_key, cursor_event_id, high_watermark_event_id, attempts,
          next_retry_epoch, lease_owner, lease_expires_epoch, last_error, created_at_epoch,
          updated_at_epoch, replay_range_id)
         VALUES (?1, ?2, ?3, ?4, ?5, ?6, 'pending', ?7, ?8, ?9, 0, NULL, NULL, NULL, NULL,
                 ?10, ?10, ?11)
         ON CONFLICT(idempotency_key) DO UPDATE SET
             status = CASE
                 WHEN extraction_tasks.status IN ('done', 'failed') THEN 'pending'
                 ELSE extraction_tasks.status
             END,
             cursor_event_id = excluded.cursor_event_id,
             high_watermark_event_id = excluded.high_watermark_event_id,
             attempts = CASE
                 WHEN extraction_tasks.status IN ('done', 'failed') THEN 0
                 ELSE extraction_tasks.attempts
             END,
             next_retry_epoch = CASE
                 WHEN extraction_tasks.status IN ('done', 'failed') THEN NULL
                 ELSE extraction_tasks.next_retry_epoch
             END,
             last_error = CASE
                 WHEN extraction_tasks.status IN ('done', 'failed') THEN NULL
                 ELSE extraction_tasks.last_error
             END,
             replay_range_id = excluded.replay_range_id,
             updated_at_epoch = excluded.updated_at_epoch",
        params![
            task_kind,
            host_id,
            workspace_id,
            project_id,
            session_row_id,
            task_kind_value.priority(),
            idempotency_key,
            from_event_id - 1,
            to_event_id,
            now,
            range_id
        ],
    )?;
    let replay_task_id: i64 = conn.query_row(
        "SELECT id FROM extraction_tasks WHERE idempotency_key = ?1",
        params![idempotency_key],
        |row| row.get(0),
    )?;
    conn.execute(
        "UPDATE extraction_replay_ranges
         SET status = 'requeued',
             replay_task_id = ?1,
             attempts = attempts + 1,
             updated_at_epoch = ?2
         WHERE id = ?3",
        params![replay_task_id, now, range_id],
    )?;
    Ok(replay_task_id)
}

#[allow(clippy::too_many_arguments)]
pub(crate) fn record_exhausted_replay_range(
    conn: &Connection,
    source_task_id: i64,
    task_kind: &str,
    host_id: i64,
    workspace_id: i64,
    project_id: i64,
    session_row_id: Option<i64>,
    from_event_id: i64,
    to_event_id: i64,
    attempts: i64,
    err: &str,
    now: i64,
) -> Result<i64> {
    conn.execute(
        "INSERT INTO extraction_replay_ranges
         (source_task_id, task_kind, host_id, workspace_id, project_id, session_row_id,
          from_event_id, to_event_id, status, attempts, last_error, created_at_epoch,
          updated_at_epoch)
         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, 'pending', ?9, ?10, ?11, ?11)
         ON CONFLICT(source_task_id, from_event_id, to_event_id) DO UPDATE SET
             status = CASE
                 WHEN extraction_replay_ranges.status = 'quarantined' THEN 'quarantined'
                 ELSE 'pending'
             END,
             attempts = excluded.attempts,
             last_error = excluded.last_error,
             updated_at_epoch = excluded.updated_at_epoch",
        params![
            source_task_id,
            task_kind,
            host_id,
            workspace_id,
            project_id,
            session_row_id,
            from_event_id,
            to_event_id,
            attempts,
            crate::db::truncate_str(err, 2000),
            now
        ],
    )?;
    conn.query_row(
        "SELECT id
         FROM extraction_replay_ranges
         WHERE source_task_id = ?1 AND from_event_id = ?2 AND to_event_id = ?3",
        params![source_task_id, from_event_id, to_event_id],
        |row| row.get(0),
    )
    .map_err(Into::into)
}

pub(crate) fn mark_replay_range_replayed_if_done(
    conn: &Connection,
    task_id: i64,
    now: i64,
) -> Result<()> {
    let updated = conn.execute(
        "UPDATE extraction_replay_ranges
         SET status = 'replayed',
             replay_task_id = COALESCE(replay_task_id, ?1),
             last_error = NULL,
             updated_at_epoch = ?2
         WHERE id = (
             SELECT replay_range_id FROM extraction_tasks
             WHERE id = ?1 AND status = 'done' AND replay_range_id IS NOT NULL
         )
           AND status != 'quarantined'
           AND NOT EXISTS (
             SELECT 1 FROM extraction_tasks t
             WHERE t.replay_range_id = extraction_replay_ranges.id
               AND t.status != 'done'
         )",
        params![task_id, now],
    )?;
    if updated > 0 {
        let range_id = conn.query_row(
            "SELECT replay_range_id
             FROM extraction_tasks
             WHERE id = ?1",
            params![task_id],
            |row| row.get::<_, i64>(0),
        )?;
        clear_terminal_failures_for_quiesced_range(conn, range_id, now)?;
    }
    Ok(())
}

fn clear_terminal_failures_for_quiesced_range(
    conn: &Connection,
    range_id: i64,
    now: i64,
) -> Result<()> {
    conn.execute(
        "UPDATE extraction_tasks
         SET status = 'done',
             attempts = 0,
             lease_owner = NULL,
             lease_expires_epoch = NULL,
             next_retry_epoch = NULL,
             last_error = NULL,
             updated_at_epoch = ?2
         WHERE replay_range_id = ?1
           AND status = 'failed'",
        params![range_id, now],
    )?;
    conn.execute(
        "UPDATE extraction_tasks
         SET status = 'done',
             attempts = 0,
             lease_owner = NULL,
             lease_expires_epoch = NULL,
             next_retry_epoch = NULL,
             last_error = NULL,
             updated_at_epoch = ?2
         WHERE id = (
             SELECT source_task_id
             FROM extraction_replay_ranges
             WHERE id = ?1
         )
           AND status = 'failed'
           AND NOT EXISTS (
             SELECT 1
             FROM extraction_replay_ranges r
             WHERE r.source_task_id = extraction_tasks.id
               AND r.status NOT IN ('replayed', 'quarantined')
         )",
        params![range_id, now],
    )?;
    Ok(())
}

pub(crate) fn mark_replay_range_failed(
    conn: &Connection,
    task_id: i64,
    now: i64,
    err: &str,
) -> Result<()> {
    conn.execute(
        "UPDATE extraction_replay_ranges
         SET status = 'failed',
             replay_task_id = COALESCE(replay_task_id, ?1),
             attempts = COALESCE((SELECT attempts FROM extraction_tasks WHERE id = ?1), attempts),
             last_error = ?2,
             updated_at_epoch = ?3
         WHERE id = (
             SELECT replay_range_id FROM extraction_tasks
             WHERE id = ?1 AND replay_range_id IS NOT NULL
         )",
        params![task_id, crate::db::truncate_str(err, 2000), now],
    )?;
    Ok(())
}
