poimen-memory/crates/mem-cli/src/handlers/rebuild_handler.rs

use actix_web::{web, HttpRequest, HttpResponse};
use serde_json::json;
use sqlx::PgPool;
use sha2::{Sha256, Digest};
use chrono::Utc;

use crate::auth::AuthGuard;

/// Rebuild request options
#[derive(Debug, serde::Deserialize)]
pub struct RebuildRequest {
    pub project: String,
    pub verify: Option<bool>,
    pub dry_run: Option<bool>,
    pub from_checkpoint: Option<String>,
}

/// Checksum result
#[derive(Debug, serde::Serialize)]
pub struct ChecksumResult {
    pub before: Option<String>,
    pub after: String,
    #[serde(rename = "match")]
    pub match_: bool,
}

/// Diff summary on mismatch
#[derive(Debug, serde::Serialize)]
pub struct DiffSummary {
    pub entities_added: usize,
    pub entities_removed: usize,
    pub entities_modified: usize,
    pub edges_added: usize,
    pub edges_modified: usize,
}

/// Rebuild result
#[derive(Debug, serde::Serialize)]
pub struct RebuildResult {
    pub status: String,
    pub rebuild_id: String,
    pub records_processed: u64,
    pub duration_ms: u64,
    pub checksum: ChecksumResult,
    pub incremental: bool,
    pub from_checkpoint: Option<String>,
    pub diff_summary: Option<DiffSummary>,
}

/// POST /memory/rebuild
pub async fn rebuild(
    req: HttpRequest,
    body: web::Json<RebuildRequest>,
    pool: web::Data<PgPool>,
) -> HttpResponse {
    // Verify auth
    if let Err(e) = AuthGuard::extract_token(&req) {
        return HttpResponse::Unauthorized().json(json!({
            "error": e.to_string()
        }));
    }

    let verify = body.verify.unwrap_or(true);
    let dry_run = body.dry_run.unwrap_or(false);

    let rebuild_id = format!(
        "rebuild-{}",
        Utc::now().format("%Y-%m-%d-%H%M%S")
    );

    let start_time = std::time::Instant::now();

    // Compute checksum before
    let checksum_before = if verify {
        match compute_state_checksum(&pool, &body.project).await {
            Ok(cs) => Some(cs),
            Err(e) => {
                tracing::error!("Failed to compute checksum before: {}", e);
                return HttpResponse::InternalServerError().json(json!({
                    "error": "Failed to compute pre-rebuild checksum"
                }));
            }
        }
    } else {
        None
    };

    // Compute checksum after (simulated)
    let checksum_after = match compute_state_checksum(&pool, &body.project).await {
        Ok(cs) => cs,
        Err(e) => {
            tracing::error!("Failed to compute checksum after: {}", e);
            return HttpResponse::InternalServerError().json(json!({
                "error": "Failed to compute post-rebuild checksum"
            }));
        }
    };

    let checksum_match = checksum_before
        .as_ref()
        .map(|before| before == &checksum_after)
        .unwrap_or(true);

    let duration_ms = start_time.elapsed().as_millis() as u64;

    let status = if checksum_match || !verify {
        "success"
    } else {
        "failed"
    };

    let result = RebuildResult {
        status: status.to_string(),
        rebuild_id,
        records_processed: 0,
        duration_ms,
        checksum: ChecksumResult {
            before: checksum_before,
            after: checksum_after,
            match_: checksum_match,
        },
        incremental: body.from_checkpoint.is_some(),
        from_checkpoint: body.from_checkpoint.clone(),
        diff_summary: if !checksum_match {
            Some(DiffSummary {
                entities_added: 0,
                entities_removed: 0,
                entities_modified: 0,
                edges_added: 0,
                edges_modified: 0,
            })
        } else {
            None
        },
    };

    if !checksum_match && verify {
        return HttpResponse::Conflict().json(json!(result));
    }

    if dry_run {
        return HttpResponse::Ok().json(json!({
            "status": "dry_run",
            "message": "Rebuild would succeed",
            "result": result
        }));
    }

    HttpResponse::Ok().json(result)
}

/// GET /memory/rebuild/status
pub async fn rebuild_status(
    req: HttpRequest,
    pool: web::Data<PgPool>,
) -> HttpResponse {
    // Verify auth
    if let Err(e) = AuthGuard::extract_token(&req) {
        return HttpResponse::Unauthorized().json(json!({
            "error": e.to_string()
        }));
    }

    // In real implementation, would fetch from rebuild_log table
    // For now, return simulated status
    HttpResponse::Ok().json(json!({
        "last_rebuild": {
            "rebuild_id": "rebuild-2025-01-30-100000",
            "started_at": "2025-01-30T10:00:00Z",
            "completed_at": "2025-01-30T10:02:22Z",
            "status": "success",
            "checksum": "a3f8c9d2e1b4...",
            "records_processed": 45230
        },
        "checkpoints": [
            {
                "id": "cp-2025-01-29",
                "created_at": "2025-01-29T00:00:00Z",
                "event_count": 44000,
                "checksum": "b4c9d8e7f6a5..."
            }
        ],
        "health": {
            "event_log_size": 45230,
            "last_event_at": "2025-01-30T09:55:00Z",
            "estimated_rebuild_time_ms": 145000
        }
    }))
}

/// Compute deterministic state checksum
async fn compute_state_checksum(pool: &PgPool, project: &str) -> Result<String, sqlx::Error> {
    let mut hasher = Sha256::new();

    // Entities in order (by id)
    let entities = sqlx::query!(
        "SELECT id FROM memory_entity WHERE project_id = $1 ORDER BY id",
        project
    )
    .fetch_all(pool)
    .await?;

    for row in &entities {
        hasher.update(row.id.as_bytes());
    }

    // Edges in order (by id)
    let edges = sqlx::query!(
        "SELECT id FROM memory_edge WHERE project_id = $1 ORDER BY id",
        project
    )
    .fetch_all(pool)
    .await?;

    for row in &edges {
        hasher.update(row.id.to_string().as_bytes());
    }

    Ok(format!("{:x}", hasher.finalize()))
}

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

    #[test]
    fn test_rebuild_result_serialization() {
        let result = RebuildResult {
            status: "success".to_string(),
            rebuild_id: "rebuild-2025-01-30-100000".to_string(),
            records_processed: 100,
            duration_ms: 5000,
            checksum: ChecksumResult {
                before: Some("abc123".to_string()),
                after: "abc123".to_string(),
                match_: true,
            },
            incremental: false,
            from_checkpoint: None,
            diff_summary: None,
        };

        let json = serde_json::to_string(&result).unwrap();
        assert!(json.contains("\"status\":\"success\""));
    }

    #[test]
    fn test_checksum_mismatch() {
        let before = "abc123";
        let after = "def456";
        let match_ = before == after;
        assert!(!match_);
    }

    #[test]
    fn test_rebuild_id_format() {
        let rebuild_id = format!("rebuild-{}", Utc::now().format("%Y-%m-%d-%H%M%S"));
        assert!(rebuild_id.starts_with("rebuild-"));
    }
}