magi-code 0.96.1

use super::store::CheckpointStore;
use crate::config::McPaths;
use serde::{Deserialize, Serialize};
use std::{
    fs,
    io::ErrorKind,
    path::{Path, PathBuf},
};

pub(crate) const CHECKPOINT_SCHEMA_VERSION: u64 = 1;
pub(crate) const MAX_SNAPSHOT_BYTES: u64 = 1024 * 1024;

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub(crate) struct SnapshotBlobRef {
    pub(crate) sha256: String,
    pub(crate) bytes: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub(crate) struct FileSnapshotEvent {
    pub(crate) session_id: String,
    pub(crate) user_turn: u64,
    pub(crate) tool: SnapshotTool,
    pub(crate) cwd: PathBuf,
    pub(crate) relative_path: PathBuf,
    pub(crate) pre: Option<SnapshotBlobRef>,
    pub(crate) post: Option<SnapshotBlobRef>,
    pub(crate) status: SnapshotCaptureStatus,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub(crate) struct SnapshotLedgerRecord {
    pub(crate) schema_version: u64,
    pub(crate) event: FileSnapshotEvent,
}

impl SnapshotLedgerRecord {
    pub(crate) fn new(event: FileSnapshotEvent) -> Self {
        Self {
            schema_version: CHECKPOINT_SCHEMA_VERSION,
            event,
        }
    }
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
#[serde(rename_all = "snake_case")]
pub(crate) enum SnapshotTool {
    WriteFile,
    Edit,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub(crate) enum SnapshotCaptureStatus {
    Captured,
    Excluded { reason: SnapshotExclusionReason },
    Oversized { bytes: u64, max_bytes: u64 },
    Unavailable { reason: String },
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub(crate) enum SnapshotExclusionReason {
    PathEscape,
    SecretPath,
    CheckpointStorage,
    SessionStorage,
    GitStorage,
    NotFile,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum SnapshotEligibility {
    Eligible,
    Excluded(SnapshotExclusionReason),
    Oversized { bytes: u64, max_bytes: u64 },
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct SnapshotContext {
    pub(crate) store: CheckpointStore,
    pub(crate) paths: McPaths,
    pub(crate) session_id: String,
    pub(crate) user_turn: u64,
}

#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub(crate) struct LedgerRead {
    pub(crate) records: Vec<SnapshotLedgerRecord>,
    pub(crate) diagnostics: Vec<String>,
}

pub(crate) fn classify_snapshot_eligibility(
    path: &Path,
    cwd_canonical: &Path,
    paths: &McPaths,
) -> SnapshotEligibility {
    if !path.starts_with(cwd_canonical) {
        return SnapshotEligibility::Excluded(SnapshotExclusionReason::PathEscape);
    }
    if path.starts_with(&paths.checkpoints) {
        return SnapshotEligibility::Excluded(SnapshotExclusionReason::CheckpointStorage);
    }
    if path.starts_with(&paths.sessions) {
        return SnapshotEligibility::Excluded(SnapshotExclusionReason::SessionStorage);
    }
    if path == paths.auth_file
        || path == paths.settings_file
        || path == paths.project_settings_file
        || paths.local_settings_file.as_deref() == Some(path)
        || secret_like_path(path)
    {
        return SnapshotEligibility::Excluded(SnapshotExclusionReason::SecretPath);
    }
    if path
        .components()
        .any(|component| component.as_os_str() == ".git")
    {
        return SnapshotEligibility::Excluded(SnapshotExclusionReason::GitStorage);
    }
    let metadata = match fs::metadata(path) {
        Ok(metadata) => metadata,
        Err(error) if error.kind() == ErrorKind::NotFound => return SnapshotEligibility::Eligible,
        Err(_) => return SnapshotEligibility::Excluded(SnapshotExclusionReason::NotFile),
    };
    if !metadata.is_file() {
        return SnapshotEligibility::Excluded(SnapshotExclusionReason::NotFile);
    }
    if metadata.len() > MAX_SNAPSHOT_BYTES {
        return SnapshotEligibility::Oversized {
            bytes: metadata.len(),
            max_bytes: MAX_SNAPSHOT_BYTES,
        };
    }
    SnapshotEligibility::Eligible
}

pub(crate) fn capture_file_snapshot(
    context: &SnapshotContext,
    cwd_canonical: &Path,
    tool: SnapshotTool,
    path: &Path,
    pre_bytes: Option<&[u8]>,
    post_bytes: Option<&[u8]>,
) -> anyhow::Result<()> {
    capture_file_snapshot_with_hook(
        context,
        cwd_canonical,
        tool,
        path,
        pre_bytes,
        post_bytes,
        || {},
    )
}

pub(super) fn capture_file_snapshot_with_hook(
    context: &SnapshotContext,
    cwd_canonical: &Path,
    tool: SnapshotTool,
    path: &Path,
    pre_bytes: Option<&[u8]>,
    post_bytes: Option<&[u8]>,
    after_blob_writes: impl FnOnce(),
) -> anyhow::Result<()> {
    let Ok(relative_path) = sanitize_relative_path(path, cwd_canonical) else {
        return Ok(());
    };
    let _mutation_lock = context.store.lock_mutations()?;
    let event = match classify_snapshot_eligibility(path, cwd_canonical, &context.paths) {
        SnapshotEligibility::Eligible => FileSnapshotEvent {
            session_id: context.session_id.clone(),
            user_turn: context.user_turn,
            tool,
            cwd: cwd_canonical.to_path_buf(),
            relative_path,
            pre: pre_bytes
                .map(|bytes| context.store.write_blob(bytes))
                .transpose()?,
            post: post_bytes
                .map(|bytes| context.store.write_blob(bytes))
                .transpose()?,
            status: SnapshotCaptureStatus::Captured,
        },
        SnapshotEligibility::Excluded(reason) => FileSnapshotEvent {
            session_id: context.session_id.clone(),
            user_turn: context.user_turn,
            tool,
            cwd: cwd_canonical.to_path_buf(),
            relative_path,
            pre: None,
            post: None,
            status: SnapshotCaptureStatus::Excluded { reason },
        },
        SnapshotEligibility::Oversized { bytes, max_bytes } => FileSnapshotEvent {
            session_id: context.session_id.clone(),
            user_turn: context.user_turn,
            tool,
            cwd: cwd_canonical.to_path_buf(),
            relative_path,
            pre: None,
            post: None,
            status: SnapshotCaptureStatus::Oversized { bytes, max_bytes },
        },
    };
    after_blob_writes();
    context
        .store
        .append_record(&SnapshotLedgerRecord::new(event))
}

pub(super) fn secret_like_path(path: &Path) -> bool {
    path.file_name()
        .and_then(|name| name.to_str())
        .is_some_and(|name| {
            name == "auth.json"
                || name == "id_rsa"
                || name == "id_ed25519"
                || name.starts_with(".env")
                || name.ends_with(".pem")
                || name.ends_with(".key")
        })
}

fn sanitize_relative_path(path: &Path, cwd: &Path) -> anyhow::Result<PathBuf> {
    let relative = path.strip_prefix(cwd).map_err(|_| {
        anyhow::anyhow!("path '{}' escapes cwd '{}'", path.display(), cwd.display())
    })?;
    if relative.as_os_str().is_empty()
        || relative
            .components()
            .any(|component| matches!(component, std::path::Component::ParentDir))
    {
        anyhow::bail!("unsafe checkpoint path");
    }
    Ok(relative.to_path_buf())
}