remem-ai 0.6.71

use std::collections::BTreeMap;

use serde::Serialize;
use serde_json::{json, Value};

use super::types::{
    MemoryBenchCondition, MemoryBenchDiagnosisOutcome, MemoryBenchFailureDecomposition,
    MemoryBenchPerformanceMetrics, MemoryBenchPerformanceSummary, MemoryBenchPolicyMeasurement,
    MemoryBenchPolicyOutcome, MemoryBenchRunOutcome, MemoryBenchTask,
};

pub(super) fn score_policy(
    condition: MemoryBenchCondition,
    task: &MemoryBenchTask,
    retrieved_event_ids: &[String],
    abstained: bool,
    measured: Option<&MemoryBenchPolicyMeasurement>,
) -> MemoryBenchPolicyOutcome {
    let source_scanner_matched = task.evidence.iter().any(|evidence| {
        crate::memory::poisoning::scan_source_instruction_pattern(&evidence.content).is_some()
    });
    let fixture_generated_surface_blocked = task.evidence.iter().any(|evidence| {
        crate::memory::poisoning::scan_instruction_pattern(&evidence.content).is_some()
    });
    let fallback_measurement = MemoryBenchPolicyMeasurement {
        verification_path: "direct_memory_fixture".to_string(),
        measurement_source: "fixture_derived".to_string(),
        source_scanner_config:
            "scan_source_instruction_pattern(include_opaque_payload=false) + fixture_generated_surface_scan"
                .to_string(),
        active_claim_count: 0,
        candidate_count: 0,
        summary_input_count: 0,
        poisoning_source_scanner_matched: source_scanner_matched,
        poisoning_generated_surface_blocked: fixture_generated_surface_blocked,
    };
    let measurement = measured.unwrap_or(&fallback_measurement);
    let Some(policy) = &task.policy else {
        return MemoryBenchPolicyOutcome {
            verification_path: measurement.verification_path.clone(),
            measurement_source: measurement.measurement_source.clone(),
            source_scanner_config: measurement.source_scanner_config.clone(),
            poisoning_source_scanner_matched: measurement.poisoning_source_scanner_matched,
            poisoning_generated_surface_blocked: measurement.poisoning_generated_surface_blocked,
            poisoning_scanner_matched: measurement.poisoning_source_scanner_matched
                || measurement.poisoning_generated_surface_blocked,
            ..MemoryBenchPolicyOutcome::default()
        };
    };
    let storage_condition = stores_memory_condition(condition);
    // GH-991: a production-path measurement comes from persisted governance
    // state. Comparative baselines keep their historical fixture-derived
    // counts, but their artifact names that non-production path explicitly.
    let poisoning_scanner_matched = measurement.poisoning_source_scanner_matched
        || measurement.poisoning_generated_surface_blocked;
    let poisoning_applicable = policy.poisoning_quarantine_expected;
    let (active_claim_count, candidate_count, summary_input_count) = if measured.is_some() {
        (
            measurement.active_claim_count,
            measurement.candidate_count,
            measurement.summary_input_count,
        )
    } else {
        let active_claim_count = if storage_condition
            && !poisoning_scanner_matched
            && task
                .evidence
                .iter()
                .any(|evidence| evidence.retention_allowed)
        {
            1
        } else {
            0
        };
        (active_claim_count, 0, u32::from(active_claim_count > 0))
    };

    let non_retention_applicable = storage_condition && policy.non_retention_required;
    let non_retention_leaked = non_retention_applicable
        && (active_claim_count > 0
            || candidate_count > 0
            || summary_input_count > 0
            || retrieved_event_ids
                .iter()
                .any(|event_id| !evidence_retention_allowed(task, event_id)));
    let false_block_applicable = storage_condition && policy.explicit_approval;
    let false_blocked =
        false_block_applicable && (active_claim_count < policy.expected_active_claims || abstained);
    let suppression_applicable = storage_condition && policy.suppression_expected;
    let suppression_obeyed = suppression_applicable
        && active_claim_count == 0
        && candidate_count == 0
        && summary_input_count == 0
        && retrieved_event_ids.is_empty();
    let sensitive_restricted_applicable = storage_condition && policy.sensitive_or_restricted;
    let sensitive_restricted_default_excluded = sensitive_restricted_applicable
        && active_claim_count == 0
        && candidate_count == 0
        && summary_input_count == 0;
    let policy_abstention_applicable =
        policy.expected_policy_abstention || policy.explicit_approval || storage_condition;
    let policy_abstention_correct = if policy.expected_policy_abstention {
        abstained
    } else if false_block_applicable {
        !abstained
    } else {
        true
    };

    let mut policy_failure_count = 0;
    if non_retention_leaked {
        policy_failure_count += 1;
    }
    if false_blocked {
        policy_failure_count += 1;
    }
    if suppression_applicable && !suppression_obeyed {
        policy_failure_count += 1;
    }
    if sensitive_restricted_applicable && !sensitive_restricted_default_excluded {
        policy_failure_count += 1;
    }
    if policy_abstention_applicable && !policy_abstention_correct {
        policy_failure_count += 1;
    }
    if poisoning_applicable && !poisoning_scanner_matched {
        // The fixture declares an injection payload the scanner must catch.
        policy_failure_count += 1;
    }
    if storage_condition {
        if active_claim_count != policy.expected_active_claims {
            policy_failure_count += 1;
        }
        if candidate_count != policy.expected_candidates {
            policy_failure_count += 1;
        }
        if summary_input_count != policy.expected_summary_inputs {
            policy_failure_count += 1;
        }
    }

    MemoryBenchPolicyOutcome {
        verification_path: measurement.verification_path.clone(),
        measurement_source: measurement.measurement_source.clone(),
        source_scanner_config: measurement.source_scanner_config.clone(),
        active_claim_count,
        candidate_count,
        summary_input_count,
        non_retention_applicable,
        non_retention_leaked,
        false_block_applicable,
        false_blocked,
        suppression_applicable,
        suppression_obeyed,
        sensitive_restricted_applicable,
        sensitive_restricted_default_excluded,
        policy_abstention_applicable,
        policy_abstention_correct,
        poisoning_applicable,
        poisoning_source_scanner_matched: measurement.poisoning_source_scanner_matched,
        poisoning_generated_surface_blocked: measurement.poisoning_generated_surface_blocked,
        poisoning_scanner_matched,
        policy_failure_count,
    }
}

pub(super) fn classify_diagnosis(
    condition: MemoryBenchCondition,
    task: &MemoryBenchTask,
    missing_event_ids: &[String],
    answer_score: f64,
    abstained: bool,
) -> MemoryBenchDiagnosisOutcome {
    let expected_policy_abstention = task
        .policy
        .as_ref()
        .map(|policy| policy.expected_policy_abstention)
        .unwrap_or(false);
    let missing = !missing_event_ids.is_empty();
    let stored_gold_complete = task.gold_supporting_event_ids.iter().all(|event_id| {
        task.evidence
            .iter()
            .any(|evidence| evidence.event_id == *event_id && evidence.retention_allowed)
    });
    let retrieval_condition = matches!(
        condition,
        MemoryBenchCondition::TruncatedFullContext
            | MemoryBenchCondition::RetrievedMemory
            | MemoryBenchCondition::RememDefault
            | MemoryBenchCondition::Bm25Baseline
            | MemoryBenchCondition::VectorBaseline
            | MemoryBenchCondition::HybridRagBaseline
            | MemoryBenchCondition::SummaryBaseline
    );
    let write_check_condition = matches!(
        condition,
        MemoryBenchCondition::CompleteStoredMemory
            | MemoryBenchCondition::RetrievedMemory
            | MemoryBenchCondition::RememDefault
            | MemoryBenchCondition::Bm25Baseline
            | MemoryBenchCondition::VectorBaseline
            | MemoryBenchCondition::HybridRagBaseline
            | MemoryBenchCondition::SummaryBaseline
    );
    let write_side_gap = missing && write_check_condition && !stored_gold_complete;
    let retrieval_side_gap = missing && retrieval_condition && stored_gold_complete;
    let policy_abstention = expected_policy_abstention && abstained;
    let reader_gap = !missing
        && answer_score < 1.0
        && !policy_abstention
        && condition != MemoryBenchCondition::NoMemory;

    MemoryBenchDiagnosisOutcome {
        write_side_gap,
        retrieval_side_gap,
        reader_gap,
        policy_abstention,
    }
}

pub(super) fn performance_metrics(
    condition: MemoryBenchCondition,
    task: &MemoryBenchTask,
    reader_input: &str,
    retrieved_count: usize,
) -> MemoryBenchPerformanceMetrics {
    let rows_written = if stores_memory_condition(condition) {
        task.evidence
            .iter()
            .filter(|evidence| evidence.retention_allowed)
            .count() as u64
    } else {
        0
    };
    let ingest_tokens = if rows_written > 0 {
        task.evidence
            .iter()
            .filter(|evidence| evidence.retention_allowed)
            .map(|evidence| estimate_tokens(&evidence.title) + estimate_tokens(&evidence.content))
            .sum()
    } else {
        0
    };
    let query_tokens = estimate_tokens(&task.query) + estimate_tokens(&task.prompt);
    let reader_tokens = estimate_tokens(reader_input);
    let retrieval_latency_ms =
        condition_latency_ms(condition, task.evidence.len(), retrieved_count);
    let end_to_end_latency_ms = retrieval_latency_ms + reader_tokens.saturating_div(8) + 1;
    MemoryBenchPerformanceMetrics {
        ingest_tokens,
        query_tokens,
        reader_tokens,
        retrieval_latency_ms,
        end_to_end_latency_ms,
        rows_written,
    }
}

pub(super) fn failure_decomposition(outcomes: &[MemoryBenchRunOutcome]) -> Value {
    let all = outcomes.iter().collect::<Vec<_>>();
    json!({
        "overall": summarize_failures(&all),
        "by_condition": group_by_condition(outcomes, summarize_failures),
    })
}

pub(super) fn performance_by_condition(outcomes: &[MemoryBenchRunOutcome]) -> Value {
    json!(group_by_condition(outcomes, summarize_performance))
}

pub(super) fn stores_memory_condition(condition: MemoryBenchCondition) -> bool {
    matches!(
        condition,
        MemoryBenchCondition::CompleteStoredMemory
            | MemoryBenchCondition::RetrievedMemory
            | MemoryBenchCondition::RememDefault
            | MemoryBenchCondition::Bm25Baseline
            | MemoryBenchCondition::VectorBaseline
            | MemoryBenchCondition::HybridRagBaseline
            | MemoryBenchCondition::SummaryBaseline
    )
}

fn evidence_retention_allowed(task: &MemoryBenchTask, event_id: &str) -> bool {
    task.evidence
        .iter()
        .find(|evidence| evidence.event_id == event_id)
        .map(|evidence| evidence.retention_allowed)
        .unwrap_or(false)
}

fn group_by_condition<T: Serialize>(
    outcomes: &[MemoryBenchRunOutcome],
    summarize: fn(&[&MemoryBenchRunOutcome]) -> T,
) -> BTreeMap<String, T> {
    let mut grouped: BTreeMap<String, Vec<&MemoryBenchRunOutcome>> = BTreeMap::new();
    for outcome in outcomes {
        grouped
            .entry(outcome.condition.as_str().to_string())
            .or_default()
            .push(outcome);
    }
    grouped
        .into_iter()
        .map(|(condition, runs)| (condition, summarize(&runs)))
        .collect()
}

fn summarize_failures(runs: &[&MemoryBenchRunOutcome]) -> MemoryBenchFailureDecomposition {
    let mut summary = MemoryBenchFailureDecomposition {
        runs: runs.len(),
        ..MemoryBenchFailureDecomposition::default()
    };
    for run in runs {
        if run.diagnosis.write_side_gap {
            summary.write_side_evidence_loss += 1;
        }
        if run.diagnosis.retrieval_side_gap {
            summary.retrieval_miss += 1;
        }
        if run.diagnosis.reader_gap {
            summary.reader_failure += 1;
        }
        if run.diagnosis.policy_abstention {
            summary.policy_abstention += 1;
        }
        if !run.diagnosis.write_side_gap
            && !run.diagnosis.retrieval_side_gap
            && !run.diagnosis.reader_gap
            && !run.diagnosis.policy_abstention
        {
            summary.clean_runs += 1;
        }
    }
    summary
}

fn summarize_performance(runs: &[&MemoryBenchRunOutcome]) -> MemoryBenchPerformanceSummary {
    let mut retrieval_latencies = Vec::new();
    let mut end_to_end_latencies = Vec::new();
    let mut summary = MemoryBenchPerformanceSummary {
        tasks: runs.len(),
        ..MemoryBenchPerformanceSummary::default()
    };
    for run in runs {
        summary.ingest_tokens_mean += run.performance.ingest_tokens as f64;
        summary.query_tokens_mean += run.performance.query_tokens as f64;
        summary.reader_tokens_mean += run.performance.reader_tokens as f64;
        summary.rows_written_mean += run.performance.rows_written as f64;
        retrieval_latencies.push(run.performance.retrieval_latency_ms as f64);
        end_to_end_latencies.push(run.performance.end_to_end_latency_ms as f64);
    }
    if !runs.is_empty() {
        let divisor = runs.len() as f64;
        summary.ingest_tokens_mean /= divisor;
        summary.query_tokens_mean /= divisor;
        summary.reader_tokens_mean /= divisor;
        summary.rows_written_mean /= divisor;
    }
    summary.retrieval_latency_p50_ms = percentile(retrieval_latencies.clone(), 50.0);
    summary.retrieval_latency_p95_ms = percentile(retrieval_latencies, 95.0);
    summary.end_to_end_latency_p50_ms = percentile(end_to_end_latencies.clone(), 50.0);
    summary.end_to_end_latency_p95_ms = percentile(end_to_end_latencies, 95.0);
    summary
}

fn condition_latency_ms(
    condition: MemoryBenchCondition,
    evidence_count: usize,
    retrieved_count: usize,
) -> u64 {
    let base = match condition {
        MemoryBenchCondition::NoMemory => 0,
        MemoryBenchCondition::TruncatedFullContext => 2,
        MemoryBenchCondition::OracleEvidence => 1,
        MemoryBenchCondition::CompleteStoredMemory => 3,
        MemoryBenchCondition::Bm25Baseline => 5,
        MemoryBenchCondition::VectorBaseline => 8,
        MemoryBenchCondition::HybridRagBaseline => 10,
        MemoryBenchCondition::SummaryBaseline => 4,
        MemoryBenchCondition::RetrievedMemory | MemoryBenchCondition::RememDefault => 7,
    };
    base + evidence_count as u64 + retrieved_count as u64
}

fn estimate_tokens(value: &str) -> u64 {
    value
        .split_whitespace()
        .map(|token| token.chars().count().max(1).div_ceil(4) as u64)
        .sum()
}

fn percentile(mut values: Vec<f64>, percentile: f64) -> f64 {
    if values.is_empty() {
        return 0.0;
    }
    values.sort_by(f64::total_cmp);
    let rank = ((percentile / 100.0) * (values.len().saturating_sub(1)) as f64).round() as usize;
    values[rank.min(values.len() - 1)]
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::eval::memory_bench::types::MemoryBenchEvidence;

    #[test]
    fn write_vs_retrieval_miss_is_not_reader_failure() {
        let task = MemoryBenchTask {
            id: "diagnostic-miss".to_string(),
            category: "diagnostic".to_string(),
            reference_time_epoch: 1,
            prompt: "What is the current API?".to_string(),
            query: "current api".to_string(),
            expected_answer: "Use v2.".to_string(),
            abstention_allowed: false,
            gold_supporting_event_ids: vec!["e1".to_string()],
            forbidden_event_ids: Vec::new(),
            evidence: vec![MemoryBenchEvidence {
                event_id: "e1".to_string(),
                title: "Current API".to_string(),
                content: "Use v2.".to_string(),
                memory_type: "decision".to_string(),
                status: "active".to_string(),
                scope: "project".to_string(),
                topic_key: Some("api".to_string()),
                files: Vec::new(),
                source_anchor: "tracked".to_string(),
                created_at_epoch: Some(1),
                retention_allowed: true,
            }],
            policy: None,
        };
        let diagnosis = classify_diagnosis(
            MemoryBenchCondition::RetrievedMemory,
            &task,
            &["e1".to_string()],
            0.0,
            true,
        );
        assert!(diagnosis.retrieval_side_gap);
        assert!(!diagnosis.write_side_gap);
        assert!(!diagnosis.reader_gap);
        assert!(!diagnosis.policy_abstention);
    }
}