//! Compliance rule registry.
//!
//! The registry is the canonical metadata for every rule_id authored under
//! `compliance/`. Corpus YAML files repeat each rule's `rule_id` and `name`
//! so the body file reads on its own, but severity, applicable transports,
//! spec versions, and references live here. The runner (future )
//! loads the registry once at startup, then joins it against the corpus
//! file bodies to decide what to evaluate.
//!
//! Why a dedicated registry: corpus files describe how to test a rule, but
//! the metadata (RFC 2119 severity, which transports a rule even applies
//! to, which spec versions claim it) is orthogonal to the assertions. Two
//! corpus files can declare the same `rule_id` (one for v2024 and one for
//! v2025), and the registry tells the runner the truth about its severity
//! once.

use std::fs;
use std::path::Path;

use serde::{Deserialize, Serialize};

use crate::{Error, Result};

/// Parsed `compliance/registry.yml` document.
///
/// Holds the version of the registry format itself plus every published
/// rule. Lookup goes through [`Self::get`] / [`Self::by_category`], so
/// callers never iterate the underlying `Vec` directly.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RuleRegistry {
    /// Format version. Bumped only when the registry's own shape changes,
    /// not when rules are added or removed.
    pub version: u32,
    /// Every rule in alphabetical order by `id`. The on-disk file is
    /// grouped by category for readability; the parser does not enforce
    /// that order at load time.
    pub rules: Vec<RuleMetadata>,
}

/// Metadata for a single compliance rule.
///
/// Fields mirror the YAML directly; serde does the heavy lifting. The
/// `references` list defaults to empty so corpus authors are not forced to
/// cite a spec section for every rule (some rules are internal robustness
/// checks with no spec hook).
///
/// The rubric fields (`section`, `applicability`, `spec_citation`) classify
/// the rule for the compliance scoring rubric. They are
/// `#[serde(default)]` so legacy files load without them, but the registry
/// validator now requires that every rule populate them.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct RuleMetadata {
    /// Stable identifier in `PREFIX-NNN` form per ADR 0015.
    pub id: String,
    /// Short human-readable summary.
    pub name: String,
    /// Category prefix the rule lives under. Matches the prefix of `id`.
    pub category: String,
    /// RFC 2119 severity.
    pub severity: RuleSeverity,
    /// Which transports the rule applies to. Typical values are
    /// `["stdio"]`, `["http"]`, or `["stdio", "http"]`.
    pub applies_to: Vec<String>,
    /// MCP spec revisions the rule targets. Free-form strings keyed on
    /// the directory names under `compliance/` (`v2024-11-05`, `draft`).
    pub spec_versions: Vec<String>,
    /// Longer description, intended for docs and reporter output.
    pub description: String,
    /// Optional list of spec sections, RFCs, or ADRs that justify the
    /// rule. Defaults to empty.
    #[serde(default)]
    pub references: Vec<String>,
    /// Rubric section bucket. One of `lifecycle`, `tools`, `resources`,
    /// `prompts`, or `transport-auth`. Independent of the
    /// `category` prefix so the rubric can group rules at a coarser
    /// grain than the ID taxonomy.
    #[serde(default)]
    pub section: String,
    /// Machine-readable predicate over the server's negotiated
    /// capabilities, evaluated by the runner to decide whether a rule
    /// applies to the server under test. Free-form CEL-like syntax for
    /// now (for example `server.capabilities.tools != null`); a future
    /// ticket will swap in a typed parser.
    #[serde(default)]
    pub applicability: String,
    /// Citation pointing back at the source of truth for the rule.
    /// Typically a spec section reference or an RFC line, distinct from
    /// the loose `references` list so reporters can render one canonical
    /// citation per check.
    #[serde(default)]
    pub spec_citation: String,
}

/// RFC 2119 severity. Stored uppercase in YAML to match the prose in the
/// registry comments and the source RFC.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "UPPERCASE")]
pub enum RuleSeverity {
    /// Hard requirement. A failure flips the compliance run red.
    Must,
    /// Strong recommendation. Surfaced as a warning by default; `--strict`
    /// promotes it to a failure.
    Should,
    /// Informational. Reporter records the outcome but never flips the
    /// exit code on its own.
    May,
}

impl RuleRegistry {
    /// Parse a registry from a YAML string.
    pub fn from_yaml(contents: &str) -> Result<Self> {
        let registry: RuleRegistry = serde_yaml::from_str(contents)?;
        registry.validate()?;
        Ok(registry)
    }

    /// Load a registry from disk.
    pub fn load(path: &Path) -> Result<Self> {
        let contents = fs::read_to_string(path)?;
        Self::from_yaml(&contents)
    }

    /// Look up a single rule by id. Returns `None` when the id is unknown.
    pub fn get(&self, rule_id: &str) -> Option<&RuleMetadata> {
        self.rules.iter().find(|r| r.id == rule_id)
    }

    /// Every rule whose `category` matches the given prefix. Useful for the
    /// runner's `--category` filter.
    pub fn by_category(&self, category: &str) -> Vec<&RuleMetadata> {
        self.rules
            .iter()
            .filter(|r| r.category == category)
            .collect()
    }

    /// Total rule count. Reporters surface this as part of the run summary.
    pub fn len(&self) -> usize {
        self.rules.len()
    }

    /// `true` when no rules are registered. Cheap shorthand for callers
    /// that want to short-circuit a missing-corpus error.
    pub fn is_empty(&self) -> bool {
        self.rules.is_empty()
    }

    /// Cross-check the registry's own invariants: rule ids are unique, the
    /// id prefix matches the declared category, and applies_to is a
    /// non-empty subset of the allowed transports.
    fn validate(&self) -> Result<()> {
        let mut seen: Vec<&str> = Vec::with_capacity(self.rules.len());
        for rule in &self.rules {
            if seen.contains(&rule.id.as_str()) {
                return Err(Error::InvalidBaseline(format!(
                    "duplicate rule id in registry: {}",
                    rule.id
                )));
            }
            seen.push(rule.id.as_str());

            let prefix = rule.id.split('-').next().unwrap_or("");
            if prefix != rule.category {
                return Err(Error::InvalidBaseline(format!(
                    "rule {} has category {} but id prefix {}",
                    rule.id, rule.category, prefix
                )));
            }

            if rule.applies_to.is_empty() {
                return Err(Error::InvalidBaseline(format!(
                    "rule {} must list at least one transport in applies_to",
                    rule.id
                )));
            }
            for transport in &rule.applies_to {
                if transport != "stdio" && transport != "http" {
                    return Err(Error::InvalidBaseline(format!(
                        "rule {} has unknown transport {:?} in applies_to",
                        rule.id, transport
                    )));
                }
            }
        }
        Ok(())
    }

    /// Count of rules grouped by rubric section.
    ///
    /// Returns a sorted vector of `(section, count)` so the rendered
    /// `stats.md` reference file is byte-stable across runs.
    pub fn section_counts(&self) -> Vec<(String, usize)> {
        let mut sections: std::collections::BTreeMap<String, usize> =
            std::collections::BTreeMap::new();
        for rule in &self.rules {
            *sections.entry(rule.section.clone()).or_insert(0) += 1;
        }
        sections.into_iter().collect()
    }

    /// Count of rules grouped by RFC 2119 severity. Output runs from
    /// highest weight (MUST) to lowest (MAY) so the rendered stats file
    /// reads top-down without re-sorting.
    pub fn severity_counts(&self) -> Vec<(RuleSeverity, usize)> {
        let mut must = 0usize;
        let mut should = 0usize;
        let mut may = 0usize;
        for rule in &self.rules {
            match rule.severity {
                RuleSeverity::Must => must += 1,
                RuleSeverity::Should => should += 1,
                RuleSeverity::May => may += 1,
            }
        }
        vec![
            (RuleSeverity::Must, must),
            (RuleSeverity::Should, should),
            (RuleSeverity::May, may),
        ]
    }

    /// Rubric-level invariants: every rule must declare a
    /// section, an applicability predicate, and a spec citation. We
    /// keep this separate from the private `validate` helper because legacy tests
    /// fabricate small in-memory registries that do not need the
    /// rubric columns. The on-disk registry passes both gates.
    pub fn validate_rubric(&self) -> Result<()> {
        const ALLOWED_SECTIONS: &[&str] = &[
            "lifecycle",
            "tools",
            "resources",
            "prompts",
            "transport-auth",
        ];
        for rule in &self.rules {
            if rule.section.is_empty() {
                return Err(Error::InvalidBaseline(format!(
                    "rule {} is missing rubric section",
                    rule.id
                )));
            }
            if !ALLOWED_SECTIONS.contains(&rule.section.as_str()) {
                return Err(Error::InvalidBaseline(format!(
                    "rule {} has unknown rubric section {:?}",
                    rule.id, rule.section
                )));
            }
            if rule.applicability.is_empty() {
                return Err(Error::InvalidBaseline(format!(
                    "rule {} is missing applicability predicate",
                    rule.id
                )));
            }
            if rule.spec_citation.is_empty() {
                return Err(Error::InvalidBaseline(format!(
                    "rule {} is missing spec_citation",
                    rule.id
                )));
            }
        }
        Ok(())
    }
}

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

    fn sample_yaml() -> &'static str {
        r#"
version: 1
rules:
  - id: PROTO-001
    name: initialize handshake required
    category: PROTO
    severity: MUST
    applies_to: [stdio, http]
    spec_versions: [v2024-11-05, draft]
    description: |
      Servers must answer initialize before other requests.
    references:
      - "MCP spec, lifecycle"
  - id: EDGE-001
    name: unicode normalization
    category: EDGE
    severity: SHOULD
    applies_to: [stdio, http]
    spec_versions: [v2024-11-05]
    description: |
      Server handles NFC and NFD identically.
"#
    }

    #[test]
    fn parses_a_well_formed_registry() {
        let reg = RuleRegistry::from_yaml(sample_yaml()).unwrap();
        assert_eq!(reg.version, 1);
        assert_eq!(reg.len(), 2);
    }

    #[test]
    fn get_returns_metadata_for_a_known_rule() {
        let reg = RuleRegistry::from_yaml(sample_yaml()).unwrap();
        let rule = reg.get("PROTO-001").unwrap();
        assert_eq!(rule.severity, RuleSeverity::Must);
        assert_eq!(rule.applies_to, vec!["stdio", "http"]);
    }

    #[test]
    fn get_returns_none_for_an_unknown_rule() {
        let reg = RuleRegistry::from_yaml(sample_yaml()).unwrap();
        assert!(reg.get("PROTO-999").is_none());
    }

    #[test]
    fn by_category_returns_only_matching_rules() {
        let reg = RuleRegistry::from_yaml(sample_yaml()).unwrap();
        let edge = reg.by_category("EDGE");
        assert_eq!(edge.len(), 1);
        assert_eq!(edge[0].id, "EDGE-001");
    }

    #[test]
    fn rejects_duplicate_rule_ids() {
        let yaml = r#"
version: 1
rules:
  - id: PROTO-001
    name: first
    category: PROTO
    severity: MUST
    applies_to: [stdio]
    spec_versions: [v2024-11-05]
    description: a
  - id: PROTO-001
    name: second
    category: PROTO
    severity: MUST
    applies_to: [stdio]
    spec_versions: [v2024-11-05]
    description: b
"#;
        let err = RuleRegistry::from_yaml(yaml).unwrap_err();
        assert!(err.to_string().contains("duplicate rule id"));
    }

    #[test]
    fn rejects_mismatched_category_prefix() {
        let yaml = r#"
version: 1
rules:
  - id: PROTO-001
    name: wrong category
    category: SCHEMA
    severity: MUST
    applies_to: [stdio]
    spec_versions: [v2024-11-05]
    description: x
"#;
        let err = RuleRegistry::from_yaml(yaml).unwrap_err();
        let msg = err.to_string();
        assert!(msg.contains("category"), "got: {msg}");
    }

    #[test]
    fn rejects_empty_applies_to() {
        let yaml = r#"
version: 1
rules:
  - id: PROTO-001
    name: empty applies_to
    category: PROTO
    severity: MUST
    applies_to: []
    spec_versions: [v2024-11-05]
    description: x
"#;
        let err = RuleRegistry::from_yaml(yaml).unwrap_err();
        assert!(err.to_string().contains("at least one transport"));
    }

    #[test]
    fn rejects_unknown_transport() {
        let yaml = r#"
version: 1
rules:
  - id: PROTO-001
    name: bad transport
    category: PROTO
    severity: MUST
    applies_to: [websocket]
    spec_versions: [v2024-11-05]
    description: x
"#;
        let err = RuleRegistry::from_yaml(yaml).unwrap_err();
        assert!(err.to_string().contains("unknown transport"));
    }

    #[test]
    fn load_reads_the_repo_registry() {
        // The corpus shipped with the workspace must always parse. This
        // doubles as a sanity check that every authored rule satisfies the
        // registry's own invariants.
        let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
            .parent()
            .and_then(|p| p.parent())
            .map(|root| root.join("compliance").join("registry.yml"))
            .expect("workspace root above crates/mcptest-core");
        let registry = RuleRegistry::load(&path).unwrap();
        assert!(
            registry.len() >= 40,
            "expected >=40 rules, got {}",
            registry.len()
        );
        assert!(registry.get("PROTO-001").is_some());
        assert!(registry.get("EDGE-009").is_some());
        assert!(registry.get("HEADER-001").is_some());
        assert!(registry.get("REPLAY-001").is_some());
        assert!(registry.get("TRANSPORT-001").is_some());
    }

    fn repo_registry() -> RuleRegistry {
        let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
            .parent()
            .and_then(|p| p.parent())
            .map(|root| root.join("compliance").join("registry.yml"))
            .expect("workspace root above crates/mcptest-core");
        RuleRegistry::load(&path).unwrap()
    }

    #[test]
    fn catalog_json_is_in_sync_with_registry() {
        // schemas/compliance-rules.json is the published, citable catalog. It is
        // generated from this registry; if it drifts, regenerate it.
        let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
            .parent()
            .and_then(|p| p.parent())
            .expect("workspace root above crates/mcptest-core");
        let registry = repo_registry();
        let expected = format!("{}\n", serde_json::to_string_pretty(&registry).unwrap());
        let committed = std::fs::read_to_string(root.join("schemas").join("compliance-rules.json"))
            .expect("schemas/compliance-rules.json exists");
        assert_eq!(
            committed, expected,
            "schemas/compliance-rules.json is stale; regenerate with: \
             cargo run -p mcptest-core --bin gen-compliance-catalog"
        );
    }

    #[test]
    fn every_rule_has_rubric_metadata_populated() {
        //  CI lint: section, applicability, spec_citation are
        // never empty in the shipped corpus.
        repo_registry().validate_rubric().unwrap();
    }

    #[test]
    fn validate_rubric_rejects_missing_section() {
        let yaml = r#"
version: 1
rules:
  - id: PROTO-001
    name: a
    category: PROTO
    severity: MUST
    applies_to: [stdio]
    spec_versions: [v2024-11-05]
    description: a
    section: ""
    applicability: "true"
    spec_citation: "MCP spec"
"#;
        let reg = RuleRegistry::from_yaml(yaml).unwrap();
        let err = reg.validate_rubric().unwrap_err();
        assert!(err.to_string().contains("rubric section"));
    }

    #[test]
    fn validate_rubric_rejects_unknown_section() {
        let yaml = r#"
version: 1
rules:
  - id: PROTO-001
    name: a
    category: PROTO
    severity: MUST
    applies_to: [stdio]
    spec_versions: [v2024-11-05]
    description: a
    section: "made-up"
    applicability: "true"
    spec_citation: "MCP spec"
"#;
        let reg = RuleRegistry::from_yaml(yaml).unwrap();
        let err = reg.validate_rubric().unwrap_err();
        assert!(err.to_string().contains("unknown rubric section"));
    }

    #[test]
    fn validate_rubric_rejects_missing_spec_citation() {
        let yaml = r#"
version: 1
rules:
  - id: PROTO-001
    name: a
    category: PROTO
    severity: MUST
    applies_to: [stdio]
    spec_versions: [v2024-11-05]
    description: a
    section: lifecycle
    applicability: "true"
"#;
        let reg = RuleRegistry::from_yaml(yaml).unwrap();
        let err = reg.validate_rubric().unwrap_err();
        assert!(err.to_string().contains("spec_citation"));
    }

    #[test]
    fn section_counts_match_committed_stats_file() {
        // The stats.md reference file is the operator-facing
        // dashboard. We pin its numbers against the live registry so a
        // corpus change that forgets to regenerate stats fails the gate.
        let registry = repo_registry();
        let mut total = 0;
        for (_section, count) in registry.section_counts() {
            total += count;
        }
        assert_eq!(total, registry.len());

        let stats = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
            .parent()
            .and_then(|p| p.parent())
            .map(|root| {
                root.join("compliance")
                    .join("rubric")
                    .join("v1")
                    .join("stats.md")
            })
            .expect("workspace root above crates/mcptest-core");
        let body = std::fs::read_to_string(&stats).expect("stats.md exists");
        let expected_total = format!("Total checks: **{}**", registry.len());
        assert!(
            body.contains(&expected_total),
            "stats.md does not match registry size; expected line {expected_total:?}"
        );
        for (section, count) in registry.section_counts() {
            let needle = format!("| {section} | {count} |");
            assert!(body.contains(&needle), "stats.md missing row {needle:?}");
        }
    }

    #[test]
    fn severity_counts_sum_to_total() {
        let registry = repo_registry();
        let total: usize = registry.severity_counts().iter().map(|(_, c)| c).sum();
        assert_eq!(total, registry.len());
    }

    #[test]
    fn rubric_checks_yaml_resolves_to_registry() {
        // `checks.yaml` is the rubric-facing entry point; we
        // keep it as a symlink to registry.yml so there is one source of
        // truth. The test loads it through the rubric path to lock the
        // contract.
        let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
            .parent()
            .and_then(|p| p.parent())
            .map(|root| {
                root.join("compliance")
                    .join("rubric")
                    .join("v1")
                    .join("checks.yaml")
            })
            .expect("workspace root above crates/mcptest-core");
        let registry = RuleRegistry::load(&path).unwrap();
        registry.validate_rubric().unwrap();
        assert!(registry.len() >= 40);
    }
}
