captchaforge 0.2.39

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;
    }

    /// Recompute `best_method` / `success_rate` / `avg_solve_time`
    /// from the current `methods` map. Used by `PatternStore::merge`
    /// after per-method values change. Pure function (no IO).
    pub(crate) fn recompute_best_method(&mut self) {
        if let Some((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()))
        {
            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 {
    // Unknown / corrupt keys used to fall back silently to
    // `SolveMethod::CrowdSourced`, which incorrectly biased
    // pattern-store reorder toward a method that may not even
    // exist in the chain. Log + fall back, so operators see the
    // drift instead of it silently shifting routing. The fallback
    // is still CrowdSourced because the chain MUST end up with a
    // valid enum value here, but the log shows the corrupt key.
    match serde_json::from_value::<SolveMethod>(serde_json::Value::String(key.to_string())) {
        Ok(m) => m,
        Err(_) => {
            tracing::warn!(
                key,
                "pattern-store method key unrecognised, falling back to CrowdSourced. \
                 This usually means a pattern file from a newer captchaforge version is \
                 being loaded by an older binary."
            );
            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.
    ///
    /// Per-method merge (NOT whole-entry replace). Each method's
    /// success/latency stats are merged independently, a malicious
    /// or buggy peer file with an inflated `sample_count` on one
    /// method can no longer impose that method as the chain's
    /// `best_method` across ALL methods. Previously a whole-entry
    /// replace ran on `other_pat.sample_count > existing.sample_count`,
    /// so a single 1B-sample-count entry won the entry wholesale.
    ///
    /// Additionally caps `sample_count` at `MAX_REASONABLE_SAMPLE_COUNT`
    /// during the merge so a hostile file with `u32::MAX` doesn't
    /// dominate the EMA forever.
    pub fn merge(&self, other: &PatternStore) {
        const MAX_REASONABLE_SAMPLE_COUNT: u32 = 1_000_000;
        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() {
            let existing = self_map.entry(key.clone()).or_insert_with(|| {
                let mut p = other_pat.clone();
                p.sample_count = p.sample_count.min(MAX_REASONABLE_SAMPLE_COUNT);
                p
            });
            // Per-method merge. Each method's `sample_count` is
            // compared independently, a peer that observed more
            // BehavioralBypass evidence wins on that method only.
            for (method_key, other_method) in &other_pat.methods {
                let other_count = other_method.sample_count.min(MAX_REASONABLE_SAMPLE_COUNT);
                existing
                    .methods
                    .entry(method_key.clone())
                    .and_modify(|m| {
                        if other_count > m.sample_count {
                            let mut copy = other_method.clone();
                            copy.sample_count = other_count;
                            *m = copy;
                        }
                    })
                    .or_insert_with(|| {
                        let mut copy = other_method.clone();
                        copy.sample_count = other_count;
                        copy
                    });
            }
            // Per-entry sample_count = sum of per-method counts so
            // an attacker can't run a side-channel via the
            // entry-level counter.
            existing.sample_count = existing
                .methods
                .values()
                .map(|m| m.sample_count)
                .sum::<u32>()
                .min(MAX_REASONABLE_SAMPLE_COUNT);
            // Recompute best_method from the merged methods so the
            // winner reflects the actual merged evidence.
            existing.recompute_best_method();
        }
    }

    /// 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)]
#[path = "pattern/tests.rs"]
mod tests;