//! Extension test-pack framework.
//!
//! The 2026-07-28 spec promotes Extensions from a side door to a
//! first-class feature: every capability beyond the core set is named by
//! a reverse-DNS id, negotiated via the `extensions` capability map, and
//! versioned independently of the spec (SEP-2133). This module owns the
//! typed model both the OSS and enterprise sides share for handling
//! that surface.
//!
//! Scope of this module:
//! - The [`ExtensionId`] type for reverse-DNS identifiers and the
//!   built-in [`KnownExtension`] enum naming the two official extensions
//!   shipped in 2026-07-28 (Tasks, MCP Apps).
//! - [`advertised_extensions`] detects extensions in a server's
//!   capability block (handed back from `initialize` or `server/discover`),
//!   so the OSS `capabilities` and `inspect` commands name them in their
//!   output even without a corresponding test pack installed.
//! - The [`ExtensionPack`] trait is the surface a test pack implements:
//!   the pack declares which extension id it handles and a `run` method
//!   the runner invokes after negotiation. The Tasks pack and
//!   the enterprise MCP Apps pack will both implement it.

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

pub mod tasks;

/// A reverse-DNS extension identifier. Per SEP-2133 the canonical form is
/// `vendor.tld/feature` (the slash separates the vendor namespace from
/// the feature name) but in capability maps the wire form is
/// `vendor.tld.feature`; we accept and round-trip the latter since it is
/// what servers actually emit.
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(transparent)]
pub struct ExtensionId(pub String);

impl ExtensionId {
    /// Build an [`ExtensionId`] from a wire-form string. No validation; we
    /// accept whatever the server hands us so the OSS surface never fails
    /// to render an advertised extension.
    pub fn new(wire: impl Into<String>) -> Self {
        Self(wire.into())
    }

    /// Borrow the wire-form id.
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl std::fmt::Display for ExtensionId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.0)
    }
}

/// The official extensions that ship in MCP 2026-07-28, plus an `Other`
/// variant for everything we do not name yet. Drives the human-readable
/// label the `capabilities` and `inspect` commands attach to a detected
/// extension and lets the enterprise tier route an id to its pack.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KnownExtension {
    /// Tasks extension (SEP-2133). Tasks graduated from experimental core
    /// to an explicit extension in 2026-07-28; this is the id servers
    /// advertise to signal they implement the new lifecycle.
    Tasks,
    /// MCP Apps extension (SEP-1865). Server-rendered HTML in a sandboxed
    /// iframe. The OSS engine recognizes the advertisement and surfaces
    /// it; the validation pack itself is enterprise.
    McpApps,
    /// An extension whose id we recognize syntactically (reverse-DNS) but
    /// which is not one of the two known categories yet.
    Other,
}

impl KnownExtension {
    /// Human-readable label shown in `capabilities` and `inspect`.
    pub fn label(self) -> &'static str {
        match self {
            KnownExtension::Tasks => "Tasks",
            KnownExtension::McpApps => "MCP Apps",
            KnownExtension::Other => "extension",
        }
    }

    /// Classify an id by its wire string. The 2026-07-28 spec does not
    /// pin a canonical wire form yet (the SEP cites `modelcontextprotocol.io`
    /// reverse-DNS but examples on the spec draft use a shorter form), so
    /// we accept either by matching on the suffix.
    pub fn classify(id: &ExtensionId) -> Self {
        let s = id.as_str();
        if s.ends_with("tasks") || s.contains(".tasks") {
            KnownExtension::Tasks
        } else if s.ends_with("apps") || s.contains(".apps") {
            KnownExtension::McpApps
        } else {
            KnownExtension::Other
        }
    }
}

/// One advertised extension detected on a server's capability block.
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct AdvertisedExtension {
    /// Reverse-DNS id the server advertised.
    pub id: ExtensionId,
    /// Best-effort classification of the id.
    #[serde(skip)]
    pub kind: KnownExtension,
    /// Optional version string from the extension's nested object. The
    /// SEP suggests `{"version": "1.0"}` as the standard shape but does
    /// not require it; we capture whatever string lives at `.version`.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub version: Option<String>,
}

impl AdvertisedExtension {
    /// Render the human-readable label the `capabilities` view uses:
    /// either `Tasks (v1.0): com.modelcontextprotocol.tasks` or, for an
    /// unrecognized id, `extension: vendor.example.thing`.
    pub fn label(&self) -> String {
        let kind = self.kind.label();
        let version = self
            .version
            .as_deref()
            .map(|v| format!(" (v{v})"))
            .unwrap_or_default();
        format!("{kind}{version}: {}", self.id)
    }
}

/// Walk a server capability block and return every advertised extension.
///
/// Looks for a top-level `extensions` object whose keys are reverse-DNS
/// extension ids. A key is considered reverse-DNS when it contains a `.`
/// (so plain capability keys like `tools` or `resources` do not light up
/// the detection). The optional nested `{"version": "..."}` is captured
/// when present.
pub fn advertised_extensions(capabilities: &Value) -> Vec<AdvertisedExtension> {
    let Some(map) = capabilities.get("extensions").and_then(Value::as_object) else {
        return Vec::new();
    };
    let mut out: Vec<AdvertisedExtension> = map
        .iter()
        .filter(|(key, _)| key.contains('.'))
        .map(|(key, value)| {
            let id = ExtensionId::new(key);
            let kind = KnownExtension::classify(&id);
            let version = value
                .get("version")
                .and_then(Value::as_str)
                .map(str::to_string);
            AdvertisedExtension { id, kind, version }
        })
        .collect();
    out.sort_by(|a, b| a.id.cmp(&b.id));
    out
}

/// Trait a test pack implements. v1 is intentionally minimal: declare the
/// extension id the pack handles, plus a hook the runner calls once the
/// pack has been negotiated. The Tasks pack and the enterprise
/// MCP Apps pack implement this.
///
/// The runner's negotiation flow:
///   1. Call `initialize` (or `server/discover`).
///   2. Walk advertised extensions via [`advertised_extensions`].
///   3. For each pack whose [`ExtensionPack::id`] matches an advertised
///      extension, invoke [`ExtensionPack::run`]; skip every pack whose
///      id is not advertised.
pub trait ExtensionPack: Send + Sync {
    /// The extension id this pack handles.
    fn id(&self) -> &ExtensionId;
    /// Pack-private version string, surfaced in reports for traceability.
    fn version(&self) -> &str;
    /// Pack run hook. Returns Ok(()) on a clean run, Err with a typed
    /// message otherwise; reporters render the failure.
    fn run(&self, context: &ExtensionContext) -> Result<(), String>;
}

/// Shared inputs every pack run sees. Held as plain data so the trait
/// implementation stays test-friendly; the runner builds this from its
/// connection and capability state.
#[derive(Debug, Clone)]
pub struct ExtensionContext {
    /// Reverse-DNS id the runner negotiated this run.
    pub id: ExtensionId,
    /// Version object the server advertised under the extension key.
    pub server_advertisement: Value,
}

/// Built-in test packs shipped with mcptest. v1 ships the Tasks pack
///; the MCP Apps pack is enterprise. New packs land
/// here as they ship; suite-level pack registration over this baseline
/// is a future ticket.
pub fn builtin_packs() -> Vec<Box<dyn ExtensionPack>> {
    vec![Box::new(tasks::TasksPack::new())]
}

/// Pair every advertised extension with every pack that handles it, in
/// the order the packs were registered. Packs whose id is not advertised
/// are skipped cleanly so a small server does not produce noisy
/// `not advertised` rows.
pub fn negotiate<'a>(
    advertised: &'a [AdvertisedExtension],
    packs: &'a [Box<dyn ExtensionPack>],
) -> Vec<(&'a AdvertisedExtension, &'a dyn ExtensionPack)> {
    packs
        .iter()
        .flat_map(|pack| {
            advertised
                .iter()
                .filter(move |adv| adv.id == *pack.id())
                .map(move |adv| (adv, pack.as_ref()))
        })
        .collect()
}

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

    #[test]
    fn classifies_tasks_and_apps_by_suffix() {
        let tasks = ExtensionId::new("io.modelcontextprotocol.tasks");
        let apps = ExtensionId::new("io.modelcontextprotocol.apps");
        let other = ExtensionId::new("com.example.widgets");
        assert_eq!(KnownExtension::classify(&tasks), KnownExtension::Tasks);
        assert_eq!(KnownExtension::classify(&apps), KnownExtension::McpApps);
        assert_eq!(KnownExtension::classify(&other), KnownExtension::Other);
    }

    #[test]
    fn empty_capabilities_advertises_nothing() {
        let caps = json!({});
        assert!(advertised_extensions(&caps).is_empty());
        let caps = json!({ "tools": { "listChanged": true } });
        assert!(advertised_extensions(&caps).is_empty());
    }

    #[test]
    fn extensions_block_is_parsed_with_versions() {
        let caps = json!({
            "tools": {},
            "extensions": {
                "io.modelcontextprotocol.tasks": { "version": "1.0" },
                "io.modelcontextprotocol.apps": {},
                "com.example.widgets": { "version": "2.3" }
            }
        });
        let detected = advertised_extensions(&caps);
        assert_eq!(detected.len(), 3);
        // Sorted by id so the report is stable.
        assert_eq!(detected[0].id.as_str(), "com.example.widgets");
        assert_eq!(detected[0].kind, KnownExtension::Other);
        assert_eq!(detected[0].version.as_deref(), Some("2.3"));
        assert_eq!(detected[1].id.as_str(), "io.modelcontextprotocol.apps");
        assert_eq!(detected[1].kind, KnownExtension::McpApps);
        assert_eq!(detected[1].version, None);
        assert_eq!(detected[2].id.as_str(), "io.modelcontextprotocol.tasks");
        assert_eq!(detected[2].kind, KnownExtension::Tasks);
    }

    #[test]
    fn non_reverse_dns_keys_under_extensions_are_skipped() {
        // A server that mistakenly puts a flat key under `extensions`
        // (which the spec does not allow) should not produce a noisy
        // entry; only reverse-DNS keys are surfaced.
        let caps = json!({
            "extensions": {
                "tasks": { "version": "0.9" }
            }
        });
        assert!(advertised_extensions(&caps).is_empty());
    }

    #[test]
    fn negotiate_pairs_matching_advertisements_with_packs() {
        let caps = json!({
            "extensions": {
                tasks::TASKS_EXTENSION_ID: { "version": "1.0" },
                "com.example.unrelated": {}
            }
        });
        let advertised = advertised_extensions(&caps);
        let packs = builtin_packs();
        let matches = negotiate(&advertised, &packs);
        assert_eq!(matches.len(), 1, "only the Tasks pack should pair");
        assert_eq!(matches[0].0.id.as_str(), tasks::TASKS_EXTENSION_ID);
        assert_eq!(matches[0].1.id().as_str(), tasks::TASKS_EXTENSION_ID);
    }

    #[test]
    fn negotiate_skips_packs_whose_extension_is_not_advertised() {
        let caps = json!({});
        let advertised = advertised_extensions(&caps);
        let packs = builtin_packs();
        let matches = negotiate(&advertised, &packs);
        assert!(matches.is_empty(), "no packs run against a smaller server");
    }

    #[test]
    fn advertised_extension_label_renders_kind_and_version() {
        let adv = AdvertisedExtension {
            id: ExtensionId::new("io.modelcontextprotocol.tasks"),
            kind: KnownExtension::Tasks,
            version: Some("1.0".into()),
        };
        let label = adv.label();
        assert!(label.starts_with("Tasks (v1.0):"));
        assert!(label.contains("io.modelcontextprotocol.tasks"));
    }
}
