use sqlx::Row;

use super::{Storage, StorageError};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ContextCheckpoint {
    pub id: String,
    pub session_id: String,
    pub generation: i64,
    pub parent_checkpoint_id: Option<String>,
    pub summary: String,
    pub first_retained_seq: i64,
    pub source_max_seq: i64,
    pub trigger_reason: String,
    pub provider_kind: String,
    pub model: String,
    pub tokens_before: i64,
    pub tokens_after: i64,
    pub degraded: bool,
    pub created_at: i64,
}

#[derive(Debug, Clone)]
pub struct NewContextCheckpoint {
    pub id: String,
    pub parent_checkpoint_id: Option<String>,
    pub summary: String,
    pub first_retained_seq: i64,
    pub source_max_seq: i64,
    pub trigger_reason: String,
    pub provider_kind: String,
    pub model: String,
    pub tokens_before: i64,
    pub tokens_after: i64,
    pub degraded: bool,
    pub created_at: i64,
}

#[derive(Debug, Clone)]
pub struct ContextCheckpointState {
    pub generation: i64,
    pub active_checkpoint_id: Option<String>,
    pub checkpoint: Option<ContextCheckpoint>,
}

impl Storage {
    pub async fn load_context_checkpoint_state(
        &self,
        session_id: &str,
    ) -> Result<ContextCheckpointState, StorageError> {
        let session = sqlx::query(
            r#"
            SELECT s.context_generation,
                   s.active_context_checkpoint_id,
                   c.id AS checkpoint_id,
                   c.session_id AS checkpoint_session_id,
                   c.generation AS checkpoint_generation,
                   c.parent_checkpoint_id,
                   c.summary,
                   c.first_retained_seq,
                   c.source_max_seq,
                   c.trigger_reason,
                   c.provider_kind,
                   c.model,
                   c.tokens_before,
                   c.tokens_after,
                   c.degraded,
                   c.created_at
            FROM sessions s
            LEFT JOIN context_checkpoints c
              ON c.id = s.active_context_checkpoint_id
             AND c.session_id = s.id
            WHERE s.id = ?
            "#,
        )
        .bind(session_id)
        .fetch_optional(self.pool())
        .await?
        .ok_or_else(|| StorageError::NotFound(session_id.to_string()))?;

        let generation = session.get("context_generation");
        let active_checkpoint_id: Option<String> = session.get("active_context_checkpoint_id");
        let checkpoint_id: Option<String> = session.get("checkpoint_id");
        let checkpoint = checkpoint_id.map(|id| ContextCheckpoint {
            id,
            session_id: session.get("checkpoint_session_id"),
            generation: session.get("checkpoint_generation"),
            parent_checkpoint_id: session.get("parent_checkpoint_id"),
            summary: session.get("summary"),
            first_retained_seq: session.get("first_retained_seq"),
            source_max_seq: session.get("source_max_seq"),
            trigger_reason: session.get("trigger_reason"),
            provider_kind: session.get("provider_kind"),
            model: session.get("model"),
            tokens_before: session.get("tokens_before"),
            tokens_after: session.get("tokens_after"),
            degraded: session.get("degraded"),
            created_at: session.get("created_at"),
        });

        Ok(ContextCheckpointState {
            generation,
            active_checkpoint_id,
            checkpoint,
        })
    }

    pub async fn commit_context_checkpoint(
        &self,
        session_id: &str,
        expected_generation: i64,
        checkpoint: &NewContextCheckpoint,
    ) -> Result<ContextCheckpoint, StorageError> {
        if checkpoint.summary.trim().is_empty() {
            return Err(StorageError::Serialization(
                "context checkpoint summary cannot be empty".to_string(),
            ));
        }
        if checkpoint.first_retained_seq < 1
            || checkpoint.source_max_seq < checkpoint.first_retained_seq
        {
            return Err(StorageError::Serialization(
                "context checkpoint sequence boundary is invalid".to_string(),
            ));
        }

        let generation = expected_generation.saturating_add(1);
        let mut tx = self.pool().begin().await?;
        let updated = sqlx::query(
            r#"
            UPDATE sessions
            SET active_context_checkpoint_id = ?,
                context_generation = context_generation + 1,
                last_compressed_message_at = ?
            WHERE id = ? AND context_generation = ? AND deleted_at IS NULL
            "#,
        )
        .bind(&checkpoint.id)
        .bind(checkpoint.created_at)
        .bind(session_id)
        .bind(expected_generation)
        .execute(&mut *tx)
        .await?;
        if updated.rows_affected() != 1 {
            tx.rollback().await?;
            return Err(StorageError::Conflict(format!(
                "stale context checkpoint generation for session {session_id}"
            )));
        }

        sqlx::query(
            r#"
            INSERT INTO context_checkpoints (
                id, session_id, generation, parent_checkpoint_id, summary,
                first_retained_seq, source_max_seq, trigger_reason,
                provider_kind, model, tokens_before, tokens_after,
                degraded, created_at
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            "#,
        )
        .bind(&checkpoint.id)
        .bind(session_id)
        .bind(generation)
        .bind(&checkpoint.parent_checkpoint_id)
        .bind(&checkpoint.summary)
        .bind(checkpoint.first_retained_seq)
        .bind(checkpoint.source_max_seq)
        .bind(&checkpoint.trigger_reason)
        .bind(&checkpoint.provider_kind)
        .bind(&checkpoint.model)
        .bind(checkpoint.tokens_before)
        .bind(checkpoint.tokens_after)
        .bind(checkpoint.degraded)
        .bind(checkpoint.created_at)
        .execute(&mut *tx)
        .await?;
        tx.commit().await?;

        Ok(ContextCheckpoint {
            id: checkpoint.id.clone(),
            session_id: session_id.to_string(),
            generation,
            parent_checkpoint_id: checkpoint.parent_checkpoint_id.clone(),
            summary: checkpoint.summary.clone(),
            first_retained_seq: checkpoint.first_retained_seq,
            source_max_seq: checkpoint.source_max_seq,
            trigger_reason: checkpoint.trigger_reason.clone(),
            provider_kind: checkpoint.provider_kind.clone(),
            model: checkpoint.model.clone(),
            tokens_before: checkpoint.tokens_before,
            tokens_after: checkpoint.tokens_after,
            degraded: checkpoint.degraded,
            created_at: checkpoint.created_at,
        })
    }

    pub async fn invalidate_context_checkpoint(
        &self,
        session_id: &str,
    ) -> Result<(), StorageError> {
        let updated = sqlx::query(
            r#"
            UPDATE sessions
            SET active_context_checkpoint_id = NULL,
                context_generation = context_generation + 1,
                last_compressed_message_at = NULL
            WHERE id = ?
            "#,
        )
        .bind(session_id)
        .execute(self.pool())
        .await?;
        if updated.rows_affected() != 1 {
            return Err(StorageError::NotFound(session_id.to_string()));
        }
        Ok(())
    }
}

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

    async fn test_storage() -> (Storage, tempfile::TempDir) {
        let directory = tempfile::tempdir().unwrap();
        let storage = Storage::new(&directory.path().join("checkpoint.db"))
            .await
            .unwrap();
        sqlx::query(
            r#"
            INSERT INTO sessions (
                id, channel, chat_id, dialog_id, title, created_at, last_active_at
            ) VALUES ('session', 'cli', 'chat', 'dialog', 'checkpoint', 1, 1)
            "#,
        )
        .execute(storage.pool())
        .await
        .unwrap();
        for seq in 1..=4 {
            sqlx::query(
                "INSERT INTO messages (id, session_id, seq, role, content, created_at) VALUES (?, 'session', ?, 'user', ?, ?)",
            )
            .bind(format!("message-{seq}"))
            .bind(seq)
            .bind(format!("message {seq}"))
            .bind(seq)
            .execute(storage.pool())
            .await
            .unwrap();
        }
        (storage, directory)
    }

    fn checkpoint(id: &str) -> NewContextCheckpoint {
        NewContextCheckpoint {
            id: id.to_string(),
            parent_checkpoint_id: None,
            summary: "durable summary".to_string(),
            first_retained_seq: 2,
            source_max_seq: 4,
            trigger_reason: "manual".to_string(),
            provider_kind: "test".to_string(),
            model: "test-model".to_string(),
            tokens_before: 100,
            tokens_after: 25,
            degraded: false,
            created_at: 10,
        }
    }

    #[tokio::test]
    async fn checkpoint_commit_is_cas_guarded_and_clear_invalidates_projection() {
        let (storage, _directory) = test_storage().await;
        let committed = storage
            .commit_context_checkpoint("session", 0, &checkpoint("cp-1"))
            .await
            .unwrap();
        assert_eq!(committed.generation, 1);

        let state = storage
            .load_context_checkpoint_state("session")
            .await
            .unwrap();
        assert_eq!(state.generation, 1);
        assert_eq!(state.active_checkpoint_id.as_deref(), Some("cp-1"));
        assert_eq!(state.checkpoint, Some(committed));
        let raw_count: i64 =
            sqlx::query_scalar("SELECT COUNT(*) FROM messages WHERE session_id = 'session'")
                .fetch_one(storage.pool())
                .await
                .unwrap();
        assert_eq!(raw_count, 4);

        sqlx::query(
            "INSERT INTO messages (id, session_id, seq, role, content, created_at) VALUES ('message-5', 'session', 5, 'assistant', 'tail', 5)",
        )
        .execute(storage.pool())
        .await
        .unwrap();
        assert_eq!(storage.get_max_message_seq("session").await.unwrap(), 5);

        let stale = storage
            .commit_context_checkpoint("session", 0, &checkpoint("cp-stale"))
            .await
            .unwrap_err();
        assert!(matches!(stale, StorageError::Conflict(_)));

        storage.clear_messages("session").await.unwrap();
        let cleared = storage
            .load_context_checkpoint_state("session")
            .await
            .unwrap();
        assert_eq!(cleared.generation, 2);
        assert!(cleared.active_checkpoint_id.is_none());
        assert!(cleared.checkpoint.is_none());

        let retained_audit_rows: i64 = sqlx::query_scalar(
            "SELECT COUNT(*) FROM context_checkpoints WHERE session_id = 'session'",
        )
        .fetch_one(storage.pool())
        .await
        .unwrap();
        assert_eq!(retained_audit_rows, 1);
    }
}
