captchaforge 0.2.4

use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap};
use std::sync::{Arc, RwLock};

use super::{CaptchaType, SolveMethod};

/// EMA smoothing for early samples (n < EMA_BOOTSTRAP_N): use α = 1/n
/// (cumulative mean) so the first few observations carry their full
/// weight. Once we have enough history, switch to a bounded α so
/// vendor migrations and detection drift register within ~10 samples
/// instead of being averaged out by hundreds of stale observations.
const EMA_BOOTSTRAP_N: u32 = 10;
const EMA_ALPHA: f32 = 0.2;

/// Per-method success tracking inside a pattern. The chain stores one
/// `CaptchaPattern` per `(domain, captcha_type)` and uses
/// `best_method()` to pick which solver to try first; `MethodStat`
/// keeps a per-method EMA so the winner can flip cleanly when a
/// previously-best method starts failing.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct MethodStat {
    /// EMA-smoothed success rate (0.0–1.0).
    pub success_rate: f32,
    /// EMA-smoothed solve time in ms.
    pub avg_solve_time: u64,
    /// How many times this method has been observed.
    pub sample_count: u32,
}

impl MethodStat {
    fn record(&mut self, success: bool, solve_time_ms: u64) {
        self.sample_count = self.sample_count.saturating_add(1);
        let alpha = if self.sample_count <= EMA_BOOTSTRAP_N {
            1.0 / self.sample_count as f32
        } else {
            EMA_ALPHA
        };
        let observation = if success { 1.0 } else { 0.0 };
        self.success_rate = self.success_rate * (1.0 - alpha) + observation * alpha;
        let prev = self.avg_solve_time as f32;
        let next = prev * (1.0 - alpha) + solve_time_ms as f32 * alpha;
        self.avg_solve_time = next.round() as u64;
    }
}

/// Per-domain knowledge about which method succeeds for a given CAPTCHA type.
///
/// Tracks per-method success rates with bounded-α EMA smoothing so
/// vendor migrations register within ~10 samples instead of being
/// drowned out by hundreds of historical observations under the older
/// cumulative-mean (α = 1/n) approach.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CaptchaPattern {
    pub domain: String,
    pub captcha_type: CaptchaType,
    pub best_method: SolveMethod,
    /// EMA-smoothed success rate of `best_method` (0.0–1.0). Kept as
    /// a top-level shortcut for backwards-compat; per-method stats
    /// live in `methods`.
    pub success_rate: f32,
    /// EMA-smoothed solve time of `best_method` in milliseconds.
    pub avg_solve_time: u64,
    /// Total attempts recorded across all methods.
    pub sample_count: u32,
    /// Per-method EMA history. Best method = highest success_rate;
    /// ties broken by fastest avg_solve_time. `BTreeMap` (not
    /// `HashMap`) so JSON serialization order is deterministic — keeps
    /// snapshot tests + cross-instance `merge` reproducible.
    #[serde(default)]
    pub methods: BTreeMap<String, MethodStat>,
}

impl CaptchaPattern {
    pub fn new(domain: impl Into<String>, captcha_type: CaptchaType, method: SolveMethod) -> Self {
        Self {
            domain: domain.into(),
            captcha_type,
            best_method: method,
            success_rate: 0.0,
            avg_solve_time: 0,
            sample_count: 0,
            methods: BTreeMap::new(),
        }
    }

    /// Update rolling stats with a new observation. `best_method` is
    /// the method with the highest EMA-smoothed `success_rate`; ties
    /// broken by fastest `avg_solve_time`.
    pub fn record(&mut self, success: bool, solve_time_ms: u64, method: SolveMethod) {
        self.sample_count = self.sample_count.saturating_add(1);
        let key = method_key(&method);
        self.methods
            .entry(key.clone())
            .or_default()
            .record(success, solve_time_ms);

        // Recompute best across all observed methods.
        let (best_key, best_stat) = self
            .methods
            .iter()
            .max_by(|(_, a), (_, b)| {
                a.success_rate
                    .partial_cmp(&b.success_rate)
                    .unwrap_or(std::cmp::Ordering::Equal)
                    .then_with(|| b.avg_solve_time.cmp(&a.avg_solve_time))
            })
            .map(|(k, s)| (k.clone(), s.clone()))
            .unwrap_or_else(|| (key, MethodStat::default()));

        self.best_method = method_from_key(&best_key);
        self.success_rate = best_stat.success_rate;
        self.avg_solve_time = best_stat.avg_solve_time;
    }
}

fn method_key(m: &SolveMethod) -> String {
    serde_json::to_value(m)
        .ok()
        .and_then(|v| v.as_str().map(|s| s.to_string()))
        .unwrap_or_else(|| format!("{:?}", m))
}

fn method_from_key(key: &str) -> SolveMethod {
    serde_json::from_value(serde_json::Value::String(key.to_string()))
        .unwrap_or(SolveMethod::CrowdSourced)
}

/// Thread-safe in-memory store of crowd-sourced patterns.
#[derive(Debug, Clone, Default)]
pub struct PatternStore {
    inner: Arc<RwLock<HashMap<String, CaptchaPattern>>>,
}

impl PatternStore {
    fn key(domain: &str, ct: &CaptchaType) -> String {
        format!("{}:{}", domain, ct)
    }

    /// Record an observation.
    pub fn record(
        &self,
        domain: &str,
        captcha_type: &CaptchaType,
        success: bool,
        time_ms: u64,
        method: SolveMethod,
    ) {
        let key = Self::key(domain, captcha_type);
        let mut map = self.inner.write().unwrap_or_else(|e| e.into_inner());
        map.entry(key)
            .and_modify(|p| p.record(success, time_ms, method.clone()))
            .or_insert_with(|| {
                let mut p = CaptchaPattern::new(domain, captcha_type.clone(), method.clone());
                p.record(success, time_ms, method);
                p
            });
    }

    /// Look up the best known method for a domain+type pair.
    pub fn best_method(&self, domain: &str, captcha_type: &CaptchaType) -> Option<SolveMethod> {
        let key = Self::key(domain, captcha_type);
        let map = self.inner.read().unwrap_or_else(|e| e.into_inner());
        map.get(&key).map(|p| p.best_method.clone())
    }

    pub fn all_patterns(&self) -> Vec<CaptchaPattern> {
        self.inner
            .read()
            .unwrap_or_else(|e| e.into_inner())
            .values()
            .cloned()
            .collect()
    }

    /// Persist every recorded pattern to a JSON file at `path`.
    ///
    /// Atomic: writes to `<path>.tmp` first then renames into place,
    /// so a concurrent reader either sees the previous version or the
    /// new one — never a half-written file.
    ///
    /// Long-running deployments call this on shutdown (and optionally
    /// after every N writes) so the next process start retains the
    /// learned routing.
    pub fn save_to_path(&self, path: impl AsRef<std::path::Path>) -> std::io::Result<()> {
        let path = path.as_ref();
        let patterns = self.all_patterns();
        let json = serde_json::to_vec_pretty(&patterns)
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
        let tmp = path.with_extension("tmp");
        std::fs::write(&tmp, json)?;
        std::fs::rename(&tmp, path)?;
        Ok(())
    }

    /// Load a [`PatternStore`] from a JSON file written by
    /// [`Self::save_to_path`]. The returned store is independent — no
    /// background sync to the file. Call `save_to_path` again to
    /// flush updates back.
    ///
    /// Missing-file errors propagate so callers can distinguish
    /// "first run, nothing to load" from "load failed for a real
    /// reason."
    pub fn load_from_path(path: impl AsRef<std::path::Path>) -> std::io::Result<Self> {
        let bytes = std::fs::read(path.as_ref())?;
        let patterns: Vec<CaptchaPattern> = serde_json::from_slice(&bytes)
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
        let mut map = HashMap::with_capacity(patterns.len());
        for pat in patterns {
            let key = Self::key(&pat.domain, &pat.captcha_type);
            map.insert(key, pat);
        }
        Ok(Self {
            inner: Arc::new(RwLock::new(map)),
        })
    }

    /// Merge another store's patterns into this one. Cross-instance
    /// share path: a worker can load patterns recorded by other
    /// workers without overwriting its own. For overlapping
    /// `(domain, captcha_type)` keys, the entry with the higher
    /// `sample_count` wins (it's seen more evidence and so its
    /// best_method estimate is more reliable).
    pub fn merge(&self, other: &PatternStore) {
        let mut self_map = self.inner.write().unwrap_or_else(|e| e.into_inner());
        let other_map = other.inner.read().unwrap_or_else(|e| e.into_inner());
        for (key, other_pat) in other_map.iter() {
            self_map
                .entry(key.clone())
                .and_modify(|existing| {
                    if other_pat.sample_count > existing.sample_count {
                        *existing = other_pat.clone();
                    }
                })
                .or_insert_with(|| other_pat.clone());
        }
    }

    /// Convenience: load from disk if the file exists, return a
    /// fresh empty store otherwise. Use this on process startup so
    /// "first ever boot" doesn't look like an error.
    pub fn load_or_default(path: impl AsRef<std::path::Path>) -> Self {
        Self::load_from_path(path).unwrap_or_default()
    }
}

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

    #[test]
    fn pattern_record_tracks_per_method_winner() {
        // EMA semantics: success_rate / avg_solve_time on the pattern
        // reflect the *best* method, not the cross-method aggregate.
        let mut p = CaptchaPattern::new(
            "example.com",
            CaptchaType::CloudflareTurnstile,
            SolveMethod::BehavioralBypass,
        );
        p.record(true, 2000, SolveMethod::BehavioralBypass);
        p.record(false, 3000, SolveMethod::VisionLLM);
        p.record(true, 1000, SolveMethod::BehavioralBypass);
        assert_eq!(p.sample_count, 3);
        assert_eq!(p.best_method, SolveMethod::BehavioralBypass);
        // BehavioralBypass: 2 samples, both success → bootstrap EMA → 1.0.
        assert!((p.success_rate - 1.0).abs() < 0.01);
        // avg_solve_time during bootstrap (n<=10): (2000*0.5)+(1000*0.5) = 1500
        assert_eq!(p.avg_solve_time, 1500);

        // Per-method stats should be tracked separately.
        assert_eq!(p.methods.len(), 2);
        let beh = p
            .methods
            .iter()
            .find(|(k, _)| k.as_str() == "behavioral_bypass")
            .map(|(_, v)| v)
            .expect("behavioral stat");
        assert_eq!(beh.sample_count, 2);
        assert!((beh.success_rate - 1.0).abs() < 0.01);
        let vlm = p
            .methods
            .iter()
            .find(|(k, _)| k.as_str() == "vision_llm")
            .map(|(_, v)| v)
            .expect("vision stat");
        assert_eq!(vlm.sample_count, 1);
        assert!(vlm.success_rate.abs() < 0.01);
    }

    #[test]
    fn ema_recency_bias_lets_best_method_flip_within_ten_samples() {
        // Burn in 50 BehavioralBypass successes — historical winner.
        let mut p = CaptchaPattern::new(
            "shifty.test",
            CaptchaType::CloudflareTurnstile,
            SolveMethod::BehavioralBypass,
        );
        for _ in 0..50 {
            p.record(true, 800, SolveMethod::BehavioralBypass);
        }
        assert_eq!(p.best_method, SolveMethod::BehavioralBypass);

        // Vendor migrates: BehavioralBypass starts always failing,
        // VisionLLM starts always succeeding. Under cumulative-mean
        // (α=1/n), Behavioral's success_rate would still be ~50/60 = 0.83
        // after 10 new samples — Vision wouldn't catch up. Under EMA
        // with α=0.2, after ~10 alternating updates Vision should be
        // the winner.
        for _ in 0..10 {
            p.record(false, 1200, SolveMethod::BehavioralBypass);
            p.record(true, 1500, SolveMethod::VisionLLM);
        }
        assert_eq!(
            p.best_method,
            SolveMethod::VisionLLM,
            "EMA must let a freshly-winning method overtake within ~10 samples; \
             behavioral_rate={}, vision_rate={}",
            p.methods
                .get("behavioral_bypass")
                .map(|s| s.success_rate)
                .unwrap_or(0.0),
            p.methods
                .get("vision_llm")
                .map(|s| s.success_rate)
                .unwrap_or(0.0),
        );
    }

    #[test]
    fn ema_bootstrap_phase_uses_cumulative_mean() {
        // First N <= EMA_BOOTSTRAP_N samples should match cumulative mean.
        let mut p = CaptchaPattern::new("x.test", CaptchaType::HCaptcha, SolveMethod::VisionLLM);
        p.record(true, 1000, SolveMethod::VisionLLM);
        p.record(true, 2000, SolveMethod::VisionLLM);
        p.record(false, 3000, SolveMethod::VisionLLM);
        // After 3 samples: cumulative success_rate = 2/3, avg = 2000.
        let stat = p.methods.get("vision_llm").unwrap();
        assert!((stat.success_rate - 2.0 / 3.0).abs() < 0.01);
        assert_eq!(stat.avg_solve_time, 2000);
    }

    #[test]
    fn pattern_store_record_and_lookup() {
        let store = PatternStore::default();
        store.record(
            "example.com",
            &CaptchaType::RecaptchaV2,
            true,
            1500,
            SolveMethod::AudioBypass,
        );
        let method = store.best_method("example.com", &CaptchaType::RecaptchaV2);
        assert_eq!(method, Some(SolveMethod::AudioBypass));
    }

    #[test]
    fn pattern_store_unknown_returns_none() {
        let store = PatternStore::default();
        assert!(store
            .best_method("unknown.com", &CaptchaType::Slider)
            .is_none());
    }

    #[test]
    fn save_then_load_round_trips_recorded_patterns() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("patterns.json");

        let original = PatternStore::default();
        original.record(
            "alpha.test",
            &CaptchaType::HCaptcha,
            true,
            1234,
            SolveMethod::VisionLLM,
        );
        original.record(
            "beta.test",
            &CaptchaType::CloudflareTurnstile,
            true,
            500,
            SolveMethod::BehavioralBypass,
        );
        original.save_to_path(&path).expect("save");

        let loaded = PatternStore::load_from_path(&path).expect("load");
        assert_eq!(
            loaded.best_method("alpha.test", &CaptchaType::HCaptcha),
            Some(SolveMethod::VisionLLM),
        );
        assert_eq!(
            loaded.best_method("beta.test", &CaptchaType::CloudflareTurnstile),
            Some(SolveMethod::BehavioralBypass),
        );
        assert_eq!(loaded.all_patterns().len(), 2);
    }

    #[test]
    fn save_to_path_is_atomic_against_concurrent_reads() {
        // Atomic semantics: the temp file is renamed into place. We
        // can't easily race in a unit test, but we can confirm the
        // .tmp file is cleaned up after a successful save.
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("patterns.json");
        let store = PatternStore::default();
        store.record(
            "x.test",
            &CaptchaType::HCaptcha,
            true,
            1000,
            SolveMethod::VisionLLM,
        );
        store.save_to_path(&path).unwrap();
        assert!(path.exists());
        assert!(
            !path.with_extension("tmp").exists(),
            "temp file should be renamed away"
        );
    }

    #[test]
    fn load_from_path_errors_on_missing_file() {
        let res = PatternStore::load_from_path("/definitely/no/such/file.json");
        assert!(res.is_err());
    }

    #[test]
    fn load_or_default_returns_empty_when_file_missing() {
        let store = PatternStore::load_or_default("/definitely/no/such/file.json");
        assert!(store.all_patterns().is_empty());
    }

    #[test]
    fn merge_keeps_higher_sample_count_winner() {
        let a = PatternStore::default();
        a.record(
            "shared.test",
            &CaptchaType::HCaptcha,
            true,
            1000,
            SolveMethod::AudioBypass,
        );

        let b = PatternStore::default();
        for _ in 0..5 {
            b.record(
                "shared.test",
                &CaptchaType::HCaptcha,
                true,
                500,
                SolveMethod::VisionLLM,
            );
        }

        a.merge(&b);
        // b had 5 samples, a had 1 — b's VisionLLM wins.
        assert_eq!(
            a.best_method("shared.test", &CaptchaType::HCaptcha),
            Some(SolveMethod::VisionLLM),
        );
    }

    #[test]
    fn merge_keeps_existing_when_self_has_more_samples() {
        let a = PatternStore::default();
        for _ in 0..10 {
            a.record(
                "shared.test",
                &CaptchaType::HCaptcha,
                true,
                500,
                SolveMethod::VisionLLM,
            );
        }
        let b = PatternStore::default();
        b.record(
            "shared.test",
            &CaptchaType::HCaptcha,
            true,
            800,
            SolveMethod::AudioBypass,
        );

        a.merge(&b);
        // a had 10 samples, b had 1 — a's VisionLLM wins.
        assert_eq!(
            a.best_method("shared.test", &CaptchaType::HCaptcha),
            Some(SolveMethod::VisionLLM),
        );
    }

    #[test]
    fn merge_inserts_disjoint_keys() {
        let a = PatternStore::default();
        a.record(
            "a.test",
            &CaptchaType::HCaptcha,
            true,
            100,
            SolveMethod::VisionLLM,
        );
        let b = PatternStore::default();
        b.record(
            "b.test",
            &CaptchaType::HCaptcha,
            true,
            100,
            SolveMethod::AudioBypass,
        );
        a.merge(&b);
        assert_eq!(a.all_patterns().len(), 2);
    }

    #[test]
    fn pattern_store_all_patterns() {
        let store = PatternStore::default();
        store.record(
            "a.com",
            &CaptchaType::HCaptcha,
            true,
            1000,
            SolveMethod::VisionLLM,
        );
        store.record(
            "b.com",
            &CaptchaType::RecaptchaV2,
            false,
            2000,
            SolveMethod::AudioBypass,
        );
        assert_eq!(store.all_patterns().len(), 2);
    }
}