scatter-proxy 0.7.0

use std::collections::HashMap;
use std::path::Path;

use serde::{Deserialize, Serialize};

use crate::metrics::ProxyHostStats;

/// Top-level persisted state written to / read from a JSON file.
#[derive(Debug, Serialize, Deserialize)]
pub struct PersistedState {
    pub version: u32,
    pub saved_at: String,
    pub proxies: HashMap<String, PersistedProxy>,
}

/// Per-proxy persisted data: its state label and per-host health stats.
#[derive(Debug, Serialize, Deserialize)]
pub struct PersistedProxy {
    pub state: String,
    pub hosts: HashMap<String, ProxyHostStats>,
}

/// Save state to a JSON file.
///
/// The write is atomic: data is first written to a temporary file next to the
/// destination, then renamed into place so readers never see a partially-written
/// file.
///
/// * `path`         – destination file path.
/// * `health_stats` – per-proxy, per-host health statistics (from `HealthTracker::get_all_stats`).
/// * `proxy_states` – per-proxy state labels (e.g. `"Active"`, `"Dead"`).
pub async fn save_state(
    path: &Path,
    health_stats: &HashMap<String, HashMap<String, ProxyHostStats>>,
    proxy_states: &HashMap<String, String>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    // Merge health_stats and proxy_states into PersistedState.
    let mut proxies: HashMap<String, PersistedProxy> = HashMap::new();

    // Seed from health_stats (the authoritative list of proxies that have data).
    for (proxy_url, hosts) in health_stats {
        let state = proxy_states
            .get(proxy_url)
            .cloned()
            .unwrap_or_else(|| "Unknown".to_string());
        proxies.insert(
            proxy_url.clone(),
            PersistedProxy {
                state,
                hosts: hosts.clone(),
            },
        );
    }

    // Also include proxies that appear in proxy_states but have no health data.
    for (proxy_url, state) in proxy_states {
        proxies
            .entry(proxy_url.clone())
            .or_insert_with(|| PersistedProxy {
                state: state.clone(),
                hosts: HashMap::new(),
            });
    }

    let persisted = PersistedState {
        version: 1,
        saved_at: chrono::Utc::now().to_rfc3339(),
        proxies,
    };

    let json = serde_json::to_string_pretty(&persisted)?;

    // Atomic write: write to a sibling temp file, then rename.
    let tmp_path = path.with_extension("tmp");
    tokio::fs::write(&tmp_path, json.as_bytes()).await?;
    tokio::fs::rename(&tmp_path, path).await?;

    Ok(())
}

/// Load persisted state from a JSON file.
///
/// Returns `Ok(None)` when the file does not exist, so callers can distinguish
/// "first run" from "corrupt file".
pub async fn load_state(
    path: &Path,
) -> Result<Option<PersistedState>, Box<dyn std::error::Error + Send + Sync>> {
    match tokio::fs::read_to_string(path).await {
        Ok(contents) => {
            let state: PersistedState = serde_json::from_str(&contents)?;
            Ok(Some(state))
        }
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
        Err(e) => Err(e.into()),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;
    use std::path::PathBuf;

    /// Helper: build a `ProxyHostStats` with the given values.
    fn make_stats(success: u32, fail: u32, rate: f64, latency: f64, consec: u32) -> ProxyHostStats {
        ProxyHostStats {
            success,
            fail,
            success_rate: rate,
            avg_latency_ms: latency,
            consecutive_fails: consec,
        }
    }

    /// Helper: create a unique temp directory for a test.
    fn temp_dir(test_name: &str) -> PathBuf {
        let dir = std::env::temp_dir().join(format!("scatter_proxy_test_{}", test_name));
        std::fs::create_dir_all(&dir).unwrap();
        dir
    }

    // -----------------------------------------------------------------------
    // PersistedState / PersistedProxy serde round-trip
    // -----------------------------------------------------------------------

    #[test]
    fn persisted_state_serializes_and_deserializes() {
        let mut hosts = HashMap::new();
        hosts.insert("example.com".to_string(), make_stats(10, 2, 0.83, 120.5, 0));

        let mut proxies = HashMap::new();
        proxies.insert(
            "socks5h://1.2.3.4:1080".to_string(),
            PersistedProxy {
                state: "Active".to_string(),
                hosts,
            },
        );

        let state = PersistedState {
            version: 1,
            saved_at: "2024-01-15T12:00:00+00:00".to_string(),
            proxies,
        };

        let json = serde_json::to_string_pretty(&state).unwrap();
        let restored: PersistedState = serde_json::from_str(&json).unwrap();

        assert_eq!(restored.version, 1);
        assert_eq!(restored.saved_at, "2024-01-15T12:00:00+00:00");
        assert_eq!(restored.proxies.len(), 1);

        let proxy = restored.proxies.get("socks5h://1.2.3.4:1080").unwrap();
        assert_eq!(proxy.state, "Active");
        assert_eq!(proxy.hosts.len(), 1);

        let h = proxy.hosts.get("example.com").unwrap();
        assert_eq!(h.success, 10);
        assert_eq!(h.fail, 2);
        assert!((h.success_rate - 0.83).abs() < 1e-9);
        assert!((h.avg_latency_ms - 120.5).abs() < 1e-9);
        assert_eq!(h.consecutive_fails, 0);
    }

    #[test]
    fn persisted_state_empty_proxies() {
        let state = PersistedState {
            version: 1,
            saved_at: "2024-01-01T00:00:00+00:00".to_string(),
            proxies: HashMap::new(),
        };

        let json = serde_json::to_string(&state).unwrap();
        let restored: PersistedState = serde_json::from_str(&json).unwrap();

        assert_eq!(restored.version, 1);
        assert!(restored.proxies.is_empty());
    }

    #[test]
    fn persisted_proxy_empty_hosts() {
        let proxy = PersistedProxy {
            state: "Dead".to_string(),
            hosts: HashMap::new(),
        };

        let json = serde_json::to_string(&proxy).unwrap();
        let restored: PersistedProxy = serde_json::from_str(&json).unwrap();

        assert_eq!(restored.state, "Dead");
        assert!(restored.hosts.is_empty());
    }

    #[test]
    fn proxy_host_stats_serde_round_trip() {
        let stats = make_stats(100, 5, 0.952, 45.3, 2);
        let json = serde_json::to_string(&stats).unwrap();
        let restored: ProxyHostStats = serde_json::from_str(&json).unwrap();

        assert_eq!(restored.success, 100);
        assert_eq!(restored.fail, 5);
        assert!((restored.success_rate - 0.952).abs() < 1e-9);
        assert!((restored.avg_latency_ms - 45.3).abs() < 1e-9);
        assert_eq!(restored.consecutive_fails, 2);
    }

    // -----------------------------------------------------------------------
    // save_state + load_state integration
    // -----------------------------------------------------------------------

    #[tokio::test]
    async fn save_and_load_round_trip() {
        let dir = temp_dir("save_load_round_trip");
        let path = dir.join("state.json");

        let mut hosts = HashMap::new();
        hosts.insert(
            "api.example.com".to_string(),
            make_stats(50, 3, 0.94, 80.0, 1),
        );

        let mut health_stats = HashMap::new();
        health_stats.insert("socks5h://10.0.0.1:1080".to_string(), hosts);

        let mut proxy_states = HashMap::new();
        proxy_states.insert("socks5h://10.0.0.1:1080".to_string(), "Active".to_string());

        save_state(&path, &health_stats, &proxy_states)
            .await
            .unwrap();

        // Verify file exists.
        assert!(path.exists());

        let loaded = load_state(&path).await.unwrap().unwrap();
        assert_eq!(loaded.version, 1);
        assert!(!loaded.saved_at.is_empty());
        assert_eq!(loaded.proxies.len(), 1);

        let proxy = loaded.proxies.get("socks5h://10.0.0.1:1080").unwrap();
        assert_eq!(proxy.state, "Active");
        assert_eq!(proxy.hosts.len(), 1);

        let h = proxy.hosts.get("api.example.com").unwrap();
        assert_eq!(h.success, 50);
        assert_eq!(h.fail, 3);

        // Clean up.
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[tokio::test]
    async fn load_state_returns_none_for_missing_file() {
        let dir = temp_dir("load_missing");
        let path = dir.join("does_not_exist.json");

        let result = load_state(&path).await.unwrap();
        assert!(result.is_none());

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[tokio::test]
    async fn load_state_returns_error_for_invalid_json() {
        let dir = temp_dir("load_invalid_json");
        let path = dir.join("bad.json");

        tokio::fs::write(&path, b"this is not valid json")
            .await
            .unwrap();

        let result = load_state(&path).await;
        assert!(result.is_err());

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[tokio::test]
    async fn save_state_is_atomic_no_tmp_left() {
        let dir = temp_dir("atomic_no_tmp");
        let path = dir.join("state.json");
        let tmp_path = dir.join("state.tmp");

        let health_stats = HashMap::new();
        let proxy_states = HashMap::new();

        save_state(&path, &health_stats, &proxy_states)
            .await
            .unwrap();

        // The temp file should have been renamed away.
        assert!(!tmp_path.exists());
        assert!(path.exists());

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[tokio::test]
    async fn save_state_overwrites_existing() {
        let dir = temp_dir("save_overwrite");
        let path = dir.join("state.json");

        // First save.
        let mut proxy_states_1 = HashMap::new();
        proxy_states_1.insert("proxy_a".to_string(), "Active".to_string());
        save_state(&path, &HashMap::new(), &proxy_states_1)
            .await
            .unwrap();

        let loaded_1 = load_state(&path).await.unwrap().unwrap();
        assert!(loaded_1.proxies.contains_key("proxy_a"));

        // Second save with different data.
        let mut proxy_states_2 = HashMap::new();
        proxy_states_2.insert("proxy_b".to_string(), "Dead".to_string());
        save_state(&path, &HashMap::new(), &proxy_states_2)
            .await
            .unwrap();

        let loaded_2 = load_state(&path).await.unwrap().unwrap();
        assert!(!loaded_2.proxies.contains_key("proxy_a"));
        assert!(loaded_2.proxies.contains_key("proxy_b"));
        assert_eq!(loaded_2.proxies.get("proxy_b").unwrap().state, "Dead");

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[tokio::test]
    async fn save_state_merges_health_and_proxy_states() {
        let dir = temp_dir("save_merge");
        let path = dir.join("state.json");

        // proxy_a has health data AND a state entry.
        // proxy_b has only a state entry (no health data).
        // proxy_c has only health data (no state entry → should default to "Unknown").

        let mut hosts_a = HashMap::new();
        hosts_a.insert("host1.com".to_string(), make_stats(10, 0, 1.0, 50.0, 0));

        let mut hosts_c = HashMap::new();
        hosts_c.insert("host2.com".to_string(), make_stats(5, 5, 0.5, 200.0, 3));

        let mut health_stats = HashMap::new();
        health_stats.insert("proxy_a".to_string(), hosts_a);
        health_stats.insert("proxy_c".to_string(), hosts_c);

        let mut proxy_states = HashMap::new();
        proxy_states.insert("proxy_a".to_string(), "Active".to_string());
        proxy_states.insert("proxy_b".to_string(), "Dead".to_string());

        save_state(&path, &health_stats, &proxy_states)
            .await
            .unwrap();

        let loaded = load_state(&path).await.unwrap().unwrap();
        assert_eq!(loaded.proxies.len(), 3);

        // proxy_a: has state "Active" and hosts.
        let pa = loaded.proxies.get("proxy_a").unwrap();
        assert_eq!(pa.state, "Active");
        assert_eq!(pa.hosts.len(), 1);
        assert!(pa.hosts.contains_key("host1.com"));

        // proxy_b: has state "Dead" and no hosts.
        let pb = loaded.proxies.get("proxy_b").unwrap();
        assert_eq!(pb.state, "Dead");
        assert!(pb.hosts.is_empty());

        // proxy_c: has health data but no explicit state → "Unknown".
        let pc = loaded.proxies.get("proxy_c").unwrap();
        assert_eq!(pc.state, "Unknown");
        assert_eq!(pc.hosts.len(), 1);
        assert!(pc.hosts.contains_key("host2.com"));

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[tokio::test]
    async fn save_state_sets_version_1() {
        let dir = temp_dir("version_check");
        let path = dir.join("state.json");

        save_state(&path, &HashMap::new(), &HashMap::new())
            .await
            .unwrap();

        let loaded = load_state(&path).await.unwrap().unwrap();
        assert_eq!(loaded.version, 1);

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[tokio::test]
    async fn save_state_sets_saved_at_iso8601() {
        let dir = temp_dir("saved_at_iso");
        let path = dir.join("state.json");

        save_state(&path, &HashMap::new(), &HashMap::new())
            .await
            .unwrap();

        let loaded = load_state(&path).await.unwrap().unwrap();

        // Should parse as a valid RFC 3339 / ISO 8601 datetime.
        let parsed = chrono::DateTime::parse_from_rfc3339(&loaded.saved_at);
        assert!(
            parsed.is_ok(),
            "saved_at is not valid RFC 3339: {}",
            loaded.saved_at
        );
    }

    #[tokio::test]
    async fn save_state_produces_pretty_json() {
        let dir = temp_dir("pretty_json");
        let path = dir.join("state.json");

        save_state(&path, &HashMap::new(), &HashMap::new())
            .await
            .unwrap();

        let raw = tokio::fs::read_to_string(&path).await.unwrap();
        // Pretty JSON should contain newlines and indentation.
        assert!(raw.contains('\n'));
        assert!(raw.contains("  "));

        let _ = std::fs::remove_dir_all(&dir);
    }

    // -----------------------------------------------------------------------
    // PersistedState Debug
    // -----------------------------------------------------------------------

    #[test]
    fn persisted_state_debug() {
        let state = PersistedState {
            version: 1,
            saved_at: "2024-01-01T00:00:00+00:00".to_string(),
            proxies: HashMap::new(),
        };
        let dbg = format!("{state:?}");
        assert!(dbg.contains("PersistedState"));
        assert!(dbg.contains("version: 1"));
    }

    #[test]
    fn persisted_proxy_debug() {
        let proxy = PersistedProxy {
            state: "Active".to_string(),
            hosts: HashMap::new(),
        };
        let dbg = format!("{proxy:?}");
        assert!(dbg.contains("PersistedProxy"));
        assert!(dbg.contains("Active"));
    }

    // -----------------------------------------------------------------------
    // Edge case: empty everything
    // -----------------------------------------------------------------------

    #[tokio::test]
    async fn save_and_load_empty_state() {
        let dir = temp_dir("empty_state");
        let path = dir.join("state.json");

        save_state(&path, &HashMap::new(), &HashMap::new())
            .await
            .unwrap();

        let loaded = load_state(&path).await.unwrap().unwrap();
        assert_eq!(loaded.version, 1);
        assert!(loaded.proxies.is_empty());

        let _ = std::fs::remove_dir_all(&dir);
    }

    // -----------------------------------------------------------------------
    // Multiple hosts per proxy
    // -----------------------------------------------------------------------

    #[tokio::test]
    async fn save_and_load_multiple_hosts_per_proxy() {
        let dir = temp_dir("multi_hosts");
        let path = dir.join("state.json");

        let mut hosts = HashMap::new();
        hosts.insert("a.com".to_string(), make_stats(1, 0, 1.0, 10.0, 0));
        hosts.insert("b.com".to_string(), make_stats(2, 1, 0.67, 20.0, 1));
        hosts.insert("c.com".to_string(), make_stats(0, 3, 0.0, 0.0, 3));

        let mut health_stats = HashMap::new();
        health_stats.insert("proxy_x".to_string(), hosts);

        let mut proxy_states = HashMap::new();
        proxy_states.insert("proxy_x".to_string(), "Active".to_string());

        save_state(&path, &health_stats, &proxy_states)
            .await
            .unwrap();

        let loaded = load_state(&path).await.unwrap().unwrap();
        let px = loaded.proxies.get("proxy_x").unwrap();
        assert_eq!(px.hosts.len(), 3);

        let hb = px.hosts.get("b.com").unwrap();
        assert_eq!(hb.success, 2);
        assert_eq!(hb.fail, 1);
        assert_eq!(hb.consecutive_fails, 1);

        let _ = std::fs::remove_dir_all(&dir);
    }
}