rprobe/

storage.rs

// File: storage.rs
// SPDX-License-Identifier: MIT OR Apache-2.0
//
// Copyright (c) 2023-2025
// - Volker Schwaberow <volker@schwaberow.de>

use anyhow::{Context, Result};
use chrono::{DateTime, Utc};
use dirs::data_local_dir;
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use sled::{Config, Db};
use std::collections::HashMap;
use std::path::{Path, PathBuf};

use crate::content_analyzer::{ContentFinding, FindingSeverity};
use crate::desync_scanner::DesyncResult;

/// Current on-disk value encoding: bincode 2 (`bincode::serde` + standard config).
pub const SCHEMA_VERSION: u32 = 2;
const SCHEMA_VERSION_KEY: &[u8] = b"schema_version";
const META_TREE: &[u8] = b"meta";

fn bincode_config() -> impl bincode::config::Config {
    bincode::config::standard()
}

fn encode_value<T: Serialize>(value: &T) -> Result<Vec<u8>> {
    bincode::serde::encode_to_vec(value, bincode_config()).context("Failed to serialize value")
}

fn decode_value_v2<T: DeserializeOwned>(bytes: &[u8]) -> Result<T> {
    let (value, _) = bincode::serde::decode_from_slice(bytes, bincode_config())
        .context("Failed to deserialize value (schema v2)")?;
    Ok(value)
}

fn decode_value_v1<T: DeserializeOwned>(bytes: &[u8]) -> Result<T> {
    bincode1::deserialize(bytes).context("Failed to deserialize value (schema v1)")
}

fn decode_value<T: DeserializeOwned>(bytes: &[u8]) -> Result<T> {
    decode_value_v2(bytes).or_else(|_| decode_value_v1(bytes))
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScanRecord {
    pub id: String,
    pub timestamp: DateTime<Utc>,
    pub url: String,
    pub status: String,
    pub detections: Vec<String>,
    pub content_findings: Vec<ContentFinding>,
    pub tls_info: HashMap<String, String>,
    pub response_time_ms: Option<u64>,
    pub response_headers: HashMap<String, String>,
    pub content_length: Option<u64>,
    pub desync_results: Vec<DesyncResult>,
    pub screenshot_path: Option<String>,
    pub robots_txt_content: Option<String>,
    pub scan_config: ScanConfig,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScanConfig {
    pub timeout: u64,
    pub http: bool,
    pub https: bool,
    pub detect_all: bool,
    pub content_analysis: bool,
    pub tls_analysis: bool,
    pub comprehensive_tls: bool,
    pub screenshot: bool,
    pub download_robots: bool,
    pub desync: bool,
    pub plugin_name: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScanSession {
    pub id: String,
    pub timestamp: DateTime<Utc>,
    pub total_urls: usize,
    pub successful_scans: usize,
    pub failed_scans: usize,
    pub duration_ms: u64,
    pub config: ScanConfig,
}

#[derive(Debug, Clone)]
pub struct HistoryQuery {
    pub url_pattern: Option<String>,
    pub start_date: Option<DateTime<Utc>>,
    pub end_date: Option<DateTime<Utc>>,
    pub min_severity: Option<FindingSeverity>,
    pub has_detections: Option<bool>,
    pub has_tls_issues: Option<bool>,
    pub has_desync_findings: Option<bool>,
    pub status_codes: Option<Vec<String>>,
    pub limit: Option<usize>,
}

impl Default for HistoryQuery {
    fn default() -> Self {
        Self {
            url_pattern: None,
            start_date: None,
            end_date: None,
            min_severity: None,
            has_detections: None,
            has_tls_issues: None,
            has_desync_findings: None,
            status_codes: None,
            limit: Some(100),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ComparisonResult {
    pub url: String,
    pub old_record: Option<ScanRecord>,
    pub new_record: Option<ScanRecord>,
    pub changes: Vec<String>,
}


#[derive(Debug, Clone)]
pub struct MigrationReport {
    pub from_version: u32,
    pub to_version: u32,
    pub scans_migrated: u64,
    pub sessions_migrated: u64,
    pub backup_path: Option<PathBuf>,
    pub dry_run: bool,
    pub skipped: bool,
}

pub struct HistoryDatabase {
    db: Db,
    db_path: PathBuf,
    scans_tree: sled::Tree,
    sessions_tree: sled::Tree,
    url_index_tree: sled::Tree,
    meta_tree: sled::Tree,
}

impl HistoryDatabase {
    pub fn resolve_path(data_dir: Option<PathBuf>) -> PathBuf {
        match data_dir {
            Some(dir) => dir.join("rprobe_history"),
            None => data_local_dir()
                .or_else(|| Some(PathBuf::from(".")))
                .unwrap()
                .join("rprobe")
                .join("history"),
        }
    }

    pub fn new(data_dir: Option<PathBuf>) -> Result<Self> {
        Self::open_at(Self::resolve_path(data_dir), true)
    }

    pub fn open_at(db_path: PathBuf, auto_migrate: bool) -> Result<Self> {
        std::fs::create_dir_all(&db_path).context("Failed to create database directory")?;

        let db = Config::default()
            .path(&db_path)
            .compression_factor(9)
            .open()
            .context("Failed to open database")?;

        let scans_tree = db
            .open_tree(b"scans")
            .context("Failed to open scans tree")?;
        let sessions_tree = db
            .open_tree(b"sessions")
            .context("Failed to open sessions tree")?;
        let url_index_tree = db
            .open_tree(b"url_index")
            .context("Failed to open URL index tree")?;
        let meta_tree = db
            .open_tree(META_TREE)
            .context("Failed to open meta tree")?;

        let mut this = Self {
            db,
            db_path,
            scans_tree,
            sessions_tree,
            url_index_tree,
            meta_tree,
        };

        if auto_migrate {
            this.ensure_current_schema(false)?;
        }

        Ok(this)
    }

    pub fn schema_version(&self) -> Result<u32> {
        match self.meta_tree.get(SCHEMA_VERSION_KEY)? {
            Some(bytes) if bytes.len() == 4 => {
                let mut arr = [0u8; 4];
                arr.copy_from_slice(&bytes);
                Ok(u32::from_le_bytes(arr))
            }
            Some(_) => anyhow::bail!("Corrupt schema_version metadata"),
            None => {
                // Legacy DBs (pre-0.10) have no marker. Empty DBs start at current.
                if self.scans_tree.is_empty() && self.sessions_tree.is_empty() {
                    Ok(SCHEMA_VERSION)
                } else {
                    Ok(1)
                }
            }
        }
    }

    fn write_schema_version(&self, version: u32) -> Result<()> {
        self.meta_tree
            .insert(SCHEMA_VERSION_KEY, &version.to_le_bytes())
            .context("Failed to write schema_version")?;
        self.db.flush().context("Failed to flush after schema_version write")?;
        Ok(())
    }

    pub fn ensure_current_schema(&mut self, dry_run: bool) -> Result<MigrationReport> {
        let current = self.schema_version()?;
        if current == SCHEMA_VERSION {
            if current != SCHEMA_VERSION {
                // unreachable
            }
            // Fresh empty DBs without meta still need the marker written.
            if self.meta_tree.get(SCHEMA_VERSION_KEY)?.is_none() {
                if !dry_run {
                    self.write_schema_version(SCHEMA_VERSION)?;
                }
            }
            return Ok(MigrationReport {
                from_version: current,
                to_version: SCHEMA_VERSION,
                scans_migrated: 0,
                sessions_migrated: 0,
                backup_path: None,
                dry_run,
                skipped: true,
            });
        }

        if current > SCHEMA_VERSION {
            anyhow::bail!(
                "Database schema version {} is newer than supported version {}",
                current,
                SCHEMA_VERSION
            );
        }

        self.migrate_v1_to_v2(dry_run)
    }

    pub fn migrate_v1_to_v2(&mut self, dry_run: bool) -> Result<MigrationReport> {
        let from_version = self.schema_version()?;
        if from_version >= SCHEMA_VERSION {
            return Ok(MigrationReport {
                from_version,
                to_version: SCHEMA_VERSION,
                scans_migrated: 0,
                sessions_migrated: 0,
                backup_path: None,
                dry_run,
                skipped: true,
            });
        }

        let backup_path = if dry_run {
            None
        } else {
            let stamp = Utc::now().format("%Y%m%d%H%M%S");
            let path = self
                .db_path
                .parent()
                .unwrap_or_else(|| Path::new("."))
                .join(format!("rprobe_history.pre-v2-backup-{}", stamp));
            self.backup_database(path.to_str().unwrap())?;
            Some(path)
        };

        let mut scans_migrated = 0u64;
        let mut sessions_migrated = 0u64;

        for result in self.scans_tree.iter() {
            let (key, value) = result.context("Failed to iterate scans during migration")?;
            let record: ScanRecord = decode_value_v1(&value)
                .with_context(|| {
                    format!(
                        "Failed to decode legacy scan {}",
                        String::from_utf8_lossy(&key)
                    )
                })?;
            if !dry_run {
                let encoded = encode_value(&record)?;
                self.scans_tree
                    .insert(key, encoded)
                    .context("Failed to rewrite migrated scan")?;
            }
            scans_migrated += 1;
        }

        for result in self.sessions_tree.iter() {
            let (key, value) = result.context("Failed to iterate sessions during migration")?;
            let session: ScanSession = decode_value_v1(&value)
                .with_context(|| {
                    format!(
                        "Failed to decode legacy session {}",
                        String::from_utf8_lossy(&key)
                    )
                })?;
            if !dry_run {
                let encoded = encode_value(&session)?;
                self.sessions_tree
                    .insert(key, encoded)
                    .context("Failed to rewrite migrated session")?;
            }
            sessions_migrated += 1;
        }

        if !dry_run {
            self.write_schema_version(SCHEMA_VERSION)?;
            self.db.flush().context("Failed to flush after migration")?;
        }

        Ok(MigrationReport {
            from_version,
            to_version: SCHEMA_VERSION,
            scans_migrated,
            sessions_migrated,
            backup_path,
            dry_run,
            skipped: false,
        })
    }

    pub fn store_scan(&self, record: &ScanRecord) -> Result<()> {
        let key = format!("{}_{}", record.timestamp.timestamp_millis(), record.id);
        let value = encode_value(record).context("Failed to serialize scan record")?;

        self.scans_tree
            .insert(key.as_bytes(), value)
            .context("Failed to store scan record")?;

        let url_key = format!(
            "{}_{}_{}",
            record.url,
            record.timestamp.timestamp_millis(),
            record.id
        );
        self.url_index_tree
            .insert(url_key.as_bytes(), key.as_bytes())
            .context("Failed to update URL index")?;

        self.db.flush().context("Failed to flush database")?;
        Ok(())
    }

    pub fn store_scans_batch(&self, records: &[ScanRecord]) -> Result<()> {
        let mut scan_batch = sled::Batch::default();
        let mut url_index_batch = sled::Batch::default();

        for record in records {
            let key = format!("{}_{}", record.timestamp.timestamp_millis(), record.id);
            let value = encode_value(record).context("Failed to serialize scan record")?;

            scan_batch.insert(key.as_bytes(), value);

            let url_key = format!(
                "{}_{}_{}",
                record.url,
                record.timestamp.timestamp_millis(),
                record.id
            );
            url_index_batch.insert(url_key.as_bytes(), key.as_bytes());
        }

        self.scans_tree
            .apply_batch(scan_batch)
            .context("Failed to store scan batch")?;

        self.url_index_tree
            .apply_batch(url_index_batch)
            .context("Failed to update URL index batch")?;

        self.db.flush().context("Failed to flush database")?;
        Ok(())
    }

    pub fn store_session(&self, session: &ScanSession) -> Result<()> {
        let key = format!("{}_{}", session.timestamp.timestamp_millis(), session.id);
        let value = encode_value(session).context("Failed to serialize scan session")?;

        self.sessions_tree
            .insert(key.as_bytes(), value)
            .context("Failed to store scan session")?;

        self.db.flush().context("Failed to flush database")?;
        Ok(())
    }

    pub fn query_scans(&self, query: &HistoryQuery) -> Result<Vec<ScanRecord>> {
        let mut results = Vec::new();

        for result in self.scans_tree.iter() {
            let (_, value) = result.context("Failed to iterate over scans")?;
            let record: ScanRecord =
                decode_value(&value).context("Failed to deserialize scan record")?;

            if self.matches_query(&record, query) {
                results.push(record);
            }

            if let Some(limit) = query.limit {
                if results.len() >= limit {
                    break;
                }
            }
        }

        results.sort_by(|a, b| b.timestamp.cmp(&a.timestamp));
        Ok(results)
    }

    pub fn get_url_history(&self, url: &str, limit: Option<usize>) -> Result<Vec<ScanRecord>> {
        let mut results = Vec::new();
        let url_prefix = format!("{}_", url);

        for result in self.url_index_tree.scan_prefix(url_prefix.as_bytes()) {
            let (_, scan_key) = result.context("Failed to scan URL index")?;

            if let Some(scan_data) = self
                .scans_tree
                .get(scan_key)
                .context("Failed to get scan from index")?
            {
                let record: ScanRecord = decode_value(&scan_data)
                    .context("Failed to deserialize indexed scan")?;
                results.push(record);
            }

            if let Some(limit) = limit {
                if results.len() >= limit {
                    break;
                }
            }
        }

        results.sort_by(|a, b| b.timestamp.cmp(&a.timestamp));
        Ok(results)
    }

    pub fn compare_scans(
        &self,
        url: &str,
        old_timestamp: DateTime<Utc>,
        new_timestamp: DateTime<Utc>,
    ) -> Result<ComparisonResult> {
        let old_record = self.find_scan_by_url_and_time(url, old_timestamp)?;
        let new_record = self.find_scan_by_url_and_time(url, new_timestamp)?;

        let changes = self.calculate_changes(&old_record, &new_record);

        Ok(ComparisonResult {
            url: url.to_string(),
            old_record,
            new_record,
            changes,
        })
    }

    pub fn get_database_stats(&self) -> Result<HashMap<String, u64>> {
        let mut stats = HashMap::new();

        stats.insert("scans_count".to_string(), self.scans_tree.len() as u64);
        stats.insert(
            "sessions_count".to_string(),
            self.sessions_tree.len() as u64,
        );
        stats.insert(
            "url_index_entries".to_string(),
            self.url_index_tree.len() as u64,
        );

        let size_on_disk = self
            .db
            .size_on_disk()
            .context("Failed to get database size")?;
        stats.insert("size_bytes".to_string(), size_on_disk);

        Ok(stats)
    }

    pub fn clean_old_data(&self, before: DateTime<Utc>) -> Result<u64> {
        let mut deleted_count = 0;
        let cutoff_timestamp = before.timestamp_millis();

        let mut keys_to_delete = Vec::new();

        for result in self.scans_tree.iter() {
            let (key, value) = result.context("Failed to iterate during cleanup")?;
            let record: ScanRecord =
                decode_value(&value).context("Failed to deserialize for cleanup")?;

            if record.timestamp.timestamp_millis() < cutoff_timestamp {
                keys_to_delete.push(key.to_vec());
            }
        }

        for key in keys_to_delete {
            self.scans_tree
                .remove(&key)
                .context("Failed to remove old scan")?;
            deleted_count += 1;
        }

        let mut session_keys_to_delete = Vec::new();
        for result in self.sessions_tree.iter() {
            let (key, value) = result.context("Failed to iterate sessions during cleanup")?;
            let session: ScanSession = decode_value(&value)
                .context("Failed to deserialize session for cleanup")?;

            if session.timestamp.timestamp_millis() < cutoff_timestamp {
                session_keys_to_delete.push(key.to_vec());
            }
        }

        for key in session_keys_to_delete {
            self.sessions_tree
                .remove(&key)
                .context("Failed to remove old session")?;
        }

        self.rebuild_url_index()?;
        self.db.flush().context("Failed to flush after cleanup")?;

        Ok(deleted_count)
    }

    pub fn compact_database(&self) -> Result<()> {
        let stats_before = self.get_database_stats()?;

        self.scans_tree
            .flush()
            .context("Failed to flush scans tree")?;
        self.sessions_tree
            .flush()
            .context("Failed to flush sessions tree")?;
        self.url_index_tree
            .flush()
            .context("Failed to flush URL index tree")?;

        let stats_after = self.get_database_stats()?;

        let size_reduction = stats_before.get("size_bytes").unwrap_or(&0)
            - stats_after.get("size_bytes").unwrap_or(&0);
        log::info!(
            "Database compacted, reduced size by {} bytes",
            size_reduction
        );

        Ok(())
    }

    pub fn backup_database(&self, backup_path: &str) -> Result<()> {
        let backup_dir = std::path::Path::new(backup_path);
        std::fs::create_dir_all(backup_dir).context("Failed to create backup directory")?;

        let backup_db = Config::default()
            .path(backup_dir)
            .open()
            .context("Failed to open backup database")?;

        let backup_scans = backup_db
            .open_tree(b"scans")
            .context("Failed to open backup scans tree")?;
        let backup_sessions = backup_db
            .open_tree(b"sessions")
            .context("Failed to open backup sessions tree")?;
        let backup_url_index = backup_db
            .open_tree(b"url_index")
            .context("Failed to open backup URL index tree")?;

        for result in self.scans_tree.iter() {
            let (key, value) = result.context("Failed to iterate scans for backup")?;
            backup_scans
                .insert(key, value)
                .context("Failed to backup scan")?;
        }

        for result in self.sessions_tree.iter() {
            let (key, value) = result.context("Failed to iterate sessions for backup")?;
            backup_sessions
                .insert(key, value)
                .context("Failed to backup session")?;
        }

        for result in self.url_index_tree.iter() {
            let (key, value) = result.context("Failed to iterate URL index for backup")?;
            backup_url_index
                .insert(key, value)
                .context("Failed to backup URL index")?;
        }

        backup_db
            .flush()
            .context("Failed to flush backup database")?;
        Ok(())
    }

    pub fn verify_database_integrity(&self) -> Result<Vec<String>> {
        let mut issues = Vec::new();

        let scans_count = self.scans_tree.len();
        let url_index_count = self.url_index_tree.len();

        if scans_count != url_index_count {
            issues.push(format!(
                "Index mismatch: {} scans vs {} URL index entries",
                scans_count, url_index_count
            ));
        }

        for result in self.scans_tree.iter() {
            let (key, value) = result.context("Failed to iterate scans for integrity check")?;

            if decode_value::<ScanRecord>(&value).is_err() {
                issues.push(format!(
                    "Corrupted scan record with key: {}",
                    String::from_utf8_lossy(&key)
                ));
            }
        }

        for result in self.sessions_tree.iter() {
            let (key, value) = result.context("Failed to iterate sessions for integrity check")?;

            if decode_value::<ScanSession>(&value).is_err() {
                issues.push(format!(
                    "Corrupted session record with key: {}",
                    String::from_utf8_lossy(&key)
                ));
            }
        }

        Ok(issues)
    }

    fn matches_query(&self, record: &ScanRecord, query: &HistoryQuery) -> bool {
        if let Some(ref pattern) = query.url_pattern {
            if !record.url.contains(pattern) {
                return false;
            }
        }

        if let Some(start) = query.start_date {
            if record.timestamp < start {
                return false;
            }
        }

        if let Some(end) = query.end_date {
            if record.timestamp > end {
                return false;
            }
        }

        if let Some(ref status_codes) = query.status_codes {
            if !status_codes.contains(&record.status) {
                return false;
            }
        }

        if let Some(has_detections) = query.has_detections {
            let has_any = !record.detections.is_empty();
            if has_detections != has_any {
                return false;
            }
        }

        if let Some(has_tls) = query.has_tls_issues {
            let has_issues =
                record.tls_info.contains_key("warnings") || record.tls_info.contains_key("errors");
            if has_tls != has_issues {
                return false;
            }
        }

        if let Some(has_desync) = query.has_desync_findings {
            let has_findings = !record.desync_results.is_empty();
            if has_desync != has_findings {
                return false;
            }
        }

        if let Some(min_severity) = &query.min_severity {
            let has_min_severity = record
                .content_findings
                .iter()
                .any(|finding| finding.severity >= *min_severity);
            if !has_min_severity {
                return false;
            }
        }

        true
    }

    fn find_scan_by_url_and_time(
        &self,
        url: &str,
        timestamp: DateTime<Utc>,
    ) -> Result<Option<ScanRecord>> {
        let timestamp_ms = timestamp.timestamp_millis();
        let url_prefix = format!("{}_", url);

        for result in self.url_index_tree.scan_prefix(url_prefix.as_bytes()) {
            let (_, scan_key) = result.context("Failed to scan for specific timestamp")?;

            if let Some(scan_data) = self
                .scans_tree
                .get(scan_key)
                .context("Failed to get scan by timestamp")?
            {
                let record: ScanRecord = decode_value(&scan_data)
                    .context("Failed to deserialize timestamped scan")?;

                let diff = (record.timestamp.timestamp_millis() - timestamp_ms).abs();
                if diff < 60_000 {
                    return Ok(Some(record));
                }
            }
        }

        Ok(None)
    }

    fn calculate_changes(&self, old: &Option<ScanRecord>, new: &Option<ScanRecord>) -> Vec<String> {
        let mut changes = Vec::new();

        match (old, new) {
            (None, Some(_)) => changes.push("New scan result".to_string()),
            (Some(_), None) => changes.push("Scan result removed".to_string()),
            (Some(old_rec), Some(new_rec)) => {
                if old_rec.status != new_rec.status {
                    changes.push(format!(
                        "Status changed: {} -> {}",
                        old_rec.status, new_rec.status
                    ));
                }

                let old_detections: std::collections::HashSet<_> =
                    old_rec.detections.iter().collect();
                let new_detections: std::collections::HashSet<_> =
                    new_rec.detections.iter().collect();

                for detection in new_detections.difference(&old_detections) {
                    changes.push(format!("New detection: {}", detection));
                }

                for detection in old_detections.difference(&new_detections) {
                    changes.push(format!("Detection removed: {}", detection));
                }

                let old_critical_findings = old_rec
                    .content_findings
                    .iter()
                    .filter(|f| f.severity == FindingSeverity::Critical)
                    .count();
                let new_critical_findings = new_rec
                    .content_findings
                    .iter()
                    .filter(|f| f.severity == FindingSeverity::Critical)
                    .count();

                if old_critical_findings != new_critical_findings {
                    changes.push(format!(
                        "Critical findings changed: {} -> {}",
                        old_critical_findings, new_critical_findings
                    ));
                }

                let old_tls_issues = old_rec.tls_info.contains_key("warnings")
                    || old_rec.tls_info.contains_key("errors");
                let new_tls_issues = new_rec.tls_info.contains_key("warnings")
                    || new_rec.tls_info.contains_key("errors");

                if old_tls_issues != new_tls_issues {
                    if new_tls_issues && !old_tls_issues {
                        changes.push("New TLS issues detected".to_string());
                    } else if !new_tls_issues && old_tls_issues {
                        changes.push("TLS issues resolved".to_string());
                    }
                }

                if old_rec.desync_results.len() != new_rec.desync_results.len() {
                    changes.push(format!(
                        "Desync results changed: {} -> {} findings",
                        old_rec.desync_results.len(),
                        new_rec.desync_results.len()
                    ));
                }

                if old_rec.response_time_ms != new_rec.response_time_ms {
                    let old_time = old_rec.response_time_ms.unwrap_or(0);
                    let new_time = new_rec.response_time_ms.unwrap_or(0);
                    let diff = (new_time as i64 - old_time as i64).abs();
                    if diff > 100 {
                        changes.push(format!(
                            "Significant response time change: {}ms -> {}ms",
                            old_time, new_time
                        ));
                    }
                }

                let old_screenshot = old_rec.screenshot_path.is_some();
                let new_screenshot = new_rec.screenshot_path.is_some();
                if old_screenshot != new_screenshot {
                    if new_screenshot {
                        changes.push("Screenshot now available".to_string());
                    } else {
                        changes.push("Screenshot no longer available".to_string());
                    }
                }
            }
            (None, None) => {}
        }

        changes
    }

    fn rebuild_url_index(&self) -> Result<()> {
        self.url_index_tree
            .clear()
            .context("Failed to clear URL index")?;

        for result in self.scans_tree.iter() {
            let (key, value) = result.context("Failed to iterate for index rebuild")?;
            let record: ScanRecord =
                decode_value(&value).context("Failed to deserialize for index rebuild")?;

            let url_key = format!(
                "{}_{}_{}",
                record.url,
                record.timestamp.timestamp_millis(),
                record.id
            );
            self.url_index_tree
                .insert(url_key.as_bytes(), &key)
                .context("Failed to rebuild URL index entry")?;
        }

        Ok(())
    }
}

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

    fn create_test_db() -> (HistoryDatabase, TempDir) {
        let temp_dir = TempDir::new().unwrap();
        let db = HistoryDatabase::new(Some(temp_dir.path().to_path_buf())).unwrap();
        (db, temp_dir)
    }

    fn create_test_record(url: &str, timestamp: DateTime<Utc>) -> ScanRecord {
        ScanRecord {
            id: uuid::Uuid::new_v4().to_string(),
            timestamp,
            url: url.to_string(),
            status: "200".to_string(),
            detections: vec!["Nginx".to_string()],
            content_findings: vec![],
            tls_info: HashMap::new(),
            response_time_ms: Some(150),
            response_headers: HashMap::new(),
            content_length: Some(1024),
            desync_results: vec![],
            screenshot_path: None,
            robots_txt_content: None,
            scan_config: ScanConfig {
                timeout: 10,
                http: true,
                https: true,
                detect_all: true,
                content_analysis: false,
                tls_analysis: false,
                comprehensive_tls: false,
                screenshot: false,
                download_robots: false,
                desync: false,
                plugin_name: None,
            },
        }
    }

    #[test]
    fn test_store_and_query_scans() {
        let (db, _temp_dir) = create_test_db();
        let now = Utc::now();
        let record = create_test_record("https://example.com", now);

        db.store_scan(&record).unwrap();

        let query = HistoryQuery {
            url_pattern: Some("example.com".to_string()),
            ..Default::default()
        };

        let results = db.query_scans(&query).unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].url, "https://example.com");
    }

    #[test]
    fn test_url_history() {
        let (db, _temp_dir) = create_test_db();
        let now = Utc::now();
        let url = "https://example.com";

        let record1 = create_test_record(url, now - chrono::Duration::days(1));
        let record2 = create_test_record(url, now);

        db.store_scan(&record1).unwrap();
        db.store_scan(&record2).unwrap();

        let history = db.get_url_history(url, Some(10)).unwrap();
        assert_eq!(history.len(), 2);
        assert!(history[0].timestamp > history[1].timestamp);
    }

    #[test]
    fn test_database_stats() {
        let (db, _temp_dir) = create_test_db();
        let now = Utc::now();
        let record = create_test_record("https://example.com", now);

        db.store_scan(&record).unwrap();

        let stats = db.get_database_stats().unwrap();
        assert_eq!(stats.get("scans_count").unwrap(), &1);
        assert!(stats.contains_key("size_bytes"));
    }

    #[test]
    fn test_clean_old_data() {
        let (db, _temp_dir) = create_test_db();
        let now = Utc::now();

        let old_record = create_test_record("https://old.com", now - chrono::Duration::days(10));
        let new_record = create_test_record("https://new.com", now);

        db.store_scan(&old_record).unwrap();
        db.store_scan(&new_record).unwrap();

        let cutoff = now - chrono::Duration::days(5);
        let deleted = db.clean_old_data(cutoff).unwrap();

        assert_eq!(deleted, 1);

        let remaining = db.query_scans(&HistoryQuery::default()).unwrap();
        assert_eq!(remaining.len(), 1);
        assert_eq!(remaining[0].url, "https://new.com");
    }

    #[test]
    fn test_batch_storage() {
        let (db, _temp_dir) = create_test_db();
        let now = Utc::now();

        let records = vec![
            create_test_record("https://example1.com", now),
            create_test_record("https://example2.com", now),
            create_test_record("https://example3.com", now),
        ];

        db.store_scans_batch(&records).unwrap();

        let stored_records = db.query_scans(&HistoryQuery::default()).unwrap();
        assert_eq!(stored_records.len(), 3);
    }

    #[test]
    fn test_response_time_storage() {
        let (db, _temp_dir) = create_test_db();
        let now = Utc::now();

        let mut record = create_test_record("https://example.com", now);
        record.response_time_ms = Some(500);

        db.store_scan(&record).unwrap();

        let stored_records = db.query_scans(&HistoryQuery::default()).unwrap();
        assert_eq!(stored_records.len(), 1);
        assert_eq!(stored_records[0].response_time_ms, Some(500));
    }

    #[test]
    fn test_desync_results_query() {
        let (db, _temp_dir) = create_test_db();
        let now = Utc::now();

        let mut record_with_desync = create_test_record("https://example.com", now);
        record_with_desync.desync_results = vec![];

        let mut record_without_desync = create_test_record("https://example2.com", now);
        record_without_desync.desync_results = vec![];

        db.store_scan(&record_with_desync).unwrap();
        db.store_scan(&record_without_desync).unwrap();

        let query = HistoryQuery {
            has_desync_findings: Some(false),
            ..Default::default()
        };

        let results = db.query_scans(&query).unwrap();
        assert_eq!(results.len(), 2);
    }

    #[test]
    fn test_database_maintenance() {
        let (db, _temp_dir) = create_test_db();
        let now = Utc::now();

        let record = create_test_record("https://example.com", now);
        db.store_scan(&record).unwrap();

        let stats_before = db.get_database_stats().unwrap();
        assert_eq!(stats_before.get("scans_count").unwrap(), &1);

        db.compact_database().unwrap();

        let stats_after = db.get_database_stats().unwrap();
        assert_eq!(stats_after.get("scans_count").unwrap(), &1);

        let integrity_issues = db.verify_database_integrity().unwrap();
        assert!(integrity_issues.is_empty());
    }

    #[test]
    fn test_scan_comparison_with_new_fields() {
        let (db, _temp_dir) = create_test_db();
        let now = Utc::now();
        let earlier = now - chrono::Duration::hours(1);

        let mut old_record = create_test_record("https://example.com", earlier);
        old_record.response_time_ms = Some(200);
        old_record.screenshot_path = None;

        let mut new_record = create_test_record("https://example.com", now);
        new_record.response_time_ms = Some(500);
        new_record.screenshot_path = Some("screenshot.png".to_string());

        db.store_scan(&old_record).unwrap();
        db.store_scan(&new_record).unwrap();

        let comparison = db
            .compare_scans("https://example.com", earlier, now)
            .unwrap();

        assert!(!comparison.changes.is_empty());
        assert!(comparison
            .changes
            .iter()
            .any(|c| c.contains("response time")));
        assert!(comparison
            .changes
            .iter()
            .any(|c| c.contains("Screenshot now available")));
    }

    #[test]
    fn test_schema_migration_v1_to_v2() {
        let (mut db, _tmp) = create_test_db();
        // Simulate legacy v1 payloads
        let record = create_test_record("https://migrate.example", Utc::now());
        let key = format!("{}_{}", record.timestamp.timestamp_millis(), record.id);
        let legacy = bincode1::serialize(&record).expect("v1 serialize");
        db.scans_tree.insert(key.as_bytes(), legacy).unwrap();
        db.meta_tree.remove(SCHEMA_VERSION_KEY).ok();
        db.db.flush().unwrap();

        assert_eq!(db.schema_version().unwrap(), 1);
        let report = db.migrate_v1_to_v2(false).unwrap();
        assert!(!report.skipped);
        assert_eq!(report.scans_migrated, 1);
        assert_eq!(db.schema_version().unwrap(), SCHEMA_VERSION);

        let loaded = db
            .query_scans(&HistoryQuery {
                url_pattern: Some("migrate.example".into()),
                ..Default::default()
            })
            .unwrap();
        assert_eq!(loaded.len(), 1);
        assert_eq!(loaded[0].url, record.url);

        let again = db.migrate_v1_to_v2(false).unwrap();
        assert!(again.skipped);
    }

}