//! Turn a catalog drift report into suite-maintenance drafts.
//!
//! `mcptest diff` tells a developer that the server changed; this module
//! tells them what to do about their suite. Given the drift report, the
//! current (post-change) catalog, and optionally the parsed suite, the
//! builder names every existing test that touches a changed or removed
//! tool and regenerates draft tests from the current schemas, using the
//! same argument synthesis and YAML shape as the coverage suggestions so
//! merged suites stay uniform. Drafts are additive snippets: nothing in
//! the plan ever rewrites or deletes an existing test.

use std::collections::{BTreeMap, BTreeSet};

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

use super::catalog_types::{ChangeKind, DiffReport};
use crate::coverage::suggest::{render_draft_with_marker, unique_name, DraftCase};
use crate::stub_gen::synthesize_args;

/// Comment placed on the first line of every regenerated draft so a
/// reader can tell diff-driven drafts from hand-authored tests (and from
/// coverage suggestions, which carry their own marker).
pub const REGEN_MARKER: &str = "# regenerated by mcptest diff";

/// Which drift category a plan item answers. Serialized in snake_case so
/// the CLI JSON output has stable wire strings.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum RegenKind {
    /// The tool exists on both sides but its schema changed.
    Changed,
    /// The tool disappeared from the current catalog.
    Removed,
    /// The tool is new in the current catalog.
    Added,
}

/// One maintenance action, carried both as structured fields (so an
/// agent can act per tool) and, when a draft exists, as the rendered
/// YAML snippet.
#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct RegenItem {
    /// Which drift category this item answers.
    pub kind: RegenKind,
    /// The affected tool.
    pub tool: String,
    /// Names of existing suite tests that call this tool. Empty when no
    /// suite was supplied or nothing references the tool.
    pub referencing_tests: Vec<String>,
    /// Human guidance when there is no draft to merge (removed tools,
    /// unsatisfiable schemas).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub note: Option<String>,
    /// Name of the regenerated draft, when one was produced.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub draft_name: Option<String>,
    /// The draft as a YAML list entry (marker comment included),
    /// indented to sit directly under a suite's `tools:` key.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub yaml: Option<String>,
}

/// The full maintenance plan for one drift report: changed tools first,
/// then removed, then added, each sorted by tool name.
#[derive(Debug, Clone, Default, PartialEq, Serialize)]
pub struct RegenPlan {
    /// One entry per affected tool, in category order.
    pub items: Vec<RegenItem>,
}

impl RegenPlan {
    /// True when the drift report left nothing to act on.
    pub fn is_empty(&self) -> bool {
        self.items.is_empty()
    }

    /// All drafts as one document under a `tools:` key, ready to paste
    /// into (or merge with) a suite. Removed tools appear only as YAML
    /// comments naming the tests to delete or rewrite, so pasting the
    /// block never deletes anything. Empty string when there is nothing
    /// to print.
    pub fn combined_yaml(&self) -> String {
        let has_drafts = self.items.iter().any(|i| i.yaml.is_some());
        if !has_drafts && !self.items.iter().any(|i| i.kind == RegenKind::Removed) {
            return String::new();
        }
        let mut out = String::new();
        if has_drafts {
            out.push_str("tools:\n");
        }
        for item in &self.items {
            if item.kind == RegenKind::Removed {
                out.push_str(&removed_comment(item));
            }
            if let Some(yaml) = &item.yaml {
                out.push_str(yaml);
            }
        }
        out
    }

    /// The wire shape `diff --format json --suggest-regen` embeds:
    /// `{"yaml": "<combined>", "items": [...]}`.
    pub fn to_json(&self) -> Value {
        json!({ "yaml": self.combined_yaml(), "items": self.items })
    }
}

/// One YAML comment line for a removed tool, naming the tests to delete
/// or rewrite. Comments keep the combined block strictly additive.
fn removed_comment(item: &RegenItem) -> String {
    if item.referencing_tests.is_empty() {
        format!(
            "  # removed tool `{}`: no suite tests reference it\n",
            item.tool
        )
    } else {
        format!(
            "  # removed tool `{}`: delete or rewrite: {}\n",
            item.tool,
            item.referencing_tests.join(", ")
        )
    }
}

/// Build the maintenance plan for `report` against the current catalog.
///
/// `current_snapshot` is the same post-change `tools/list` value handed
/// to the diff, so drafts regenerate from the schemas the server serves
/// now. `suite` is the parsed suite document (JSON form of the YAML);
/// without it there is nothing to connect changed or removed tools to,
/// so the plan carries added-tool drafts only and the caller should say
/// so. `server_name` fills each draft's `server:` field.
///
/// ```
/// use mcptest_core::diff::{build_regen_plan, diff_tool_catalogs};
/// use serde_json::json;
///
/// let old = json!({ "tools": [] });
/// let new = json!({ "tools": [{ "name": "ping", "inputSchema": { "type": "object" } }] });
/// let report = diff_tool_catalogs(&old, &new).unwrap();
/// let plan = build_regen_plan(&report, &new, None, "server");
/// assert!(plan.combined_yaml().contains("# regenerated by mcptest diff"));
/// ```
pub fn build_regen_plan(
    report: &DiffReport,
    current_snapshot: &Value,
    suite: Option<&Value>,
    server_name: &str,
) -> RegenPlan {
    let catalog = current_tools(current_snapshot);
    let refs = suite.map(suite_references);
    let mut used: BTreeSet<String> = refs
        .as_ref()
        .map(|(_, names)| names.clone())
        .unwrap_or_default();
    let by_tool = refs.map(|(by_tool, _)| by_tool).unwrap_or_default();
    let (changed, removed, added) = categorize(report);

    let mut items = Vec::new();
    if suite.is_some() {
        for tool in &changed {
            items.push(changed_item(
                tool,
                &catalog,
                &by_tool,
                server_name,
                &mut used,
            ));
        }
        for tool in &removed {
            items.push(RegenItem {
                kind: RegenKind::Removed,
                tool: tool.clone(),
                referencing_tests: tests_for(&by_tool, tool),
                note: Some("tool removed; delete or rewrite the referencing tests".to_string()),
                draft_name: None,
                yaml: None,
            });
        }
    }
    for tool in &added {
        items.push(added_item(tool, &catalog, &by_tool, server_name, &mut used));
    }
    RegenPlan { items }
}

/// Split the report's changes into the three tool sets the plan acts
/// on. "Changed" means an input or output schema delta; description and
/// annotation edits do not invalidate a test's arguments, so they stay
/// in the drift report only.
fn categorize(report: &DiffReport) -> (BTreeSet<String>, BTreeSet<String>, BTreeSet<String>) {
    let mut changed = BTreeSet::new();
    let mut removed = BTreeSet::new();
    let mut added = BTreeSet::new();
    for change in &report.changes {
        match &change.kind {
            ChangeKind::InputSchemaChanged { .. } | ChangeKind::OutputSchemaChanged { .. } => {
                changed.insert(change.tool.clone());
            }
            ChangeKind::Removed => {
                removed.insert(change.tool.clone());
            }
            ChangeKind::Added => {
                added.insert(change.tool.clone());
            }
            // Description, annotation, title, server-instructions, and
            // response-shape drift do not change a tool's *argument* shape, so
            // they need no regenerated test and stay in the drift report only.
            // A response-shape change is about the tool's output, not the
            // inputs a draft would supply.
            ChangeKind::DescriptionChanged { .. }
            | ChangeKind::AnnotationChanged { .. }
            | ChangeKind::TitleChanged { .. }
            | ChangeKind::InstructionsChanged { .. }
            | ChangeKind::ResponseShapeChanged { .. } => {}
        }
    }
    (changed, removed, added)
}

/// Plan entry for a schema-changed tool: the referencing test names plus
/// a draft regenerated from the current schema. An unsatisfiable schema
/// yields a note instead of a draft that could never pass.
fn changed_item(
    tool: &str,
    catalog: &BTreeMap<String, Value>,
    by_tool: &BTreeMap<String, Vec<String>>,
    server_name: &str,
    used: &mut BTreeSet<String>,
) -> RegenItem {
    let referencing_tests = tests_for(by_tool, tool);
    let base = format!("{tool}: regenerated valid arguments");
    match draft_for(tool, catalog, &base, server_name, used) {
        Some((draft_name, yaml)) => RegenItem {
            kind: RegenKind::Changed,
            tool: tool.to_owned(),
            referencing_tests,
            note: Some(
                "schema changed; review the referencing tests against the regenerated draft"
                    .to_string(),
            ),
            draft_name: Some(draft_name),
            yaml: Some(yaml),
        },
        None => RegenItem {
            kind: RegenKind::Changed,
            tool: tool.to_owned(),
            referencing_tests,
            note: Some(
                "schema changed but no draft could be synthesized; rewrite by hand".to_string(),
            ),
            draft_name: None,
            yaml: None,
        },
    }
}

/// Plan entry for an added tool: a fresh draft in the same shape (and
/// with the same base name) as a coverage suggestion for an uncovered
/// tool, so the two paths produce interchangeable snippets.
fn added_item(
    tool: &str,
    catalog: &BTreeMap<String, Value>,
    by_tool: &BTreeMap<String, Vec<String>>,
    server_name: &str,
    used: &mut BTreeSet<String>,
) -> RegenItem {
    let base = format!("{tool}: valid arguments");
    match draft_for(tool, catalog, &base, server_name, used) {
        Some((draft_name, yaml)) => RegenItem {
            kind: RegenKind::Added,
            tool: tool.to_owned(),
            referencing_tests: tests_for(by_tool, tool),
            note: None,
            draft_name: Some(draft_name),
            yaml: Some(yaml),
        },
        None => RegenItem {
            kind: RegenKind::Added,
            tool: tool.to_owned(),
            referencing_tests: tests_for(by_tool, tool),
            note: Some(
                "tool added but no draft could be synthesized; write a test by hand".to_string(),
            ),
            draft_name: None,
            yaml: None,
        },
    }
}

/// Synthesize and render one happy-path draft for `tool` from the
/// current catalog. `None` when the tool is absent from the snapshot or
/// its schema is unsatisfiable, mirroring what coverage suggestions do.
fn draft_for(
    tool: &str,
    catalog: &BTreeMap<String, Value>,
    base_name: &str,
    server_name: &str,
    used: &mut BTreeSet<String>,
) -> Option<(String, String)> {
    let schema = input_schema_of(catalog.get(tool)?);
    let happy = synthesize_args(&schema);
    let args: Map<String, Value> = happy.value.as_ref().and_then(Value::as_object)?.clone();
    let name = unique_name(base_name, used);
    let yaml = render_draft_with_marker(
        &DraftCase {
            name: &name,
            server: server_name,
            tool,
            inject_error: None,
            args: &args,
            expect_target: "result",
            expect_message: "result should be a structured tool result",
        },
        REGEN_MARKER,
    );
    Some((name, yaml))
}

/// Tool catalog from the current snapshot, keyed by name. Accepts the
/// same envelope shapes the diff accepts (`tools` at the top level or
/// under `result`); anything else yields an empty map, which the plan
/// surfaces as "no draft could be synthesized" notes.
fn current_tools(snapshot: &Value) -> BTreeMap<String, Value> {
    let tools = snapshot.get("tools").and_then(Value::as_array).or_else(|| {
        snapshot
            .get("result")
            .and_then(|r| r.get("tools"))
            .and_then(Value::as_array)
    });
    let mut out = BTreeMap::new();
    for tool in tools.into_iter().flatten() {
        if let Some(name) = tool.get("name").and_then(Value::as_str) {
            out.insert(name.to_owned(), tool.clone());
        }
    }
    out
}

/// Input schema in either wire (`inputSchema`) or snake_case spelling,
/// matching what the generators and coverage suggestions accept.
fn input_schema_of(tool: &Value) -> Value {
    tool.get("inputSchema")
        .or_else(|| tool.get("input_schema"))
        .cloned()
        .unwrap_or(Value::Null)
}

/// Walk the suite's `tools[]` entries and index test names by the tool
/// they call, plus the flat set of all names for draft dedup. Entries
/// without a name or tool are skipped; the loader rejects those shapes
/// anyway.
fn suite_references(suite: &Value) -> (BTreeMap<String, Vec<String>>, BTreeSet<String>) {
    let mut by_tool: BTreeMap<String, Vec<String>> = BTreeMap::new();
    let mut names = BTreeSet::new();
    let tests = suite.get("tools").and_then(Value::as_array);
    for test in tests.into_iter().flatten() {
        let Some(name) = test.get("name").and_then(Value::as_str) else {
            continue;
        };
        names.insert(name.to_owned());
        if let Some(tool) = test.get("tool").and_then(Value::as_str) {
            by_tool
                .entry(tool.to_owned())
                .or_default()
                .push(name.to_owned());
        }
    }
    (by_tool, names)
}

/// Referencing test names for one tool, empty when nothing calls it.
fn tests_for(by_tool: &BTreeMap<String, Vec<String>>, tool: &str) -> Vec<String> {
    by_tool.get(tool).cloned().unwrap_or_default()
}

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

    /// Baseline catalog: a tool whose enum will change, a tool that will
    /// be removed, and a bystander that never changes.
    fn baseline() -> Value {
        json!({ "tools": [
            {
                "name": "create_issue",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "title": { "type": "string" },
                        "priority": { "type": "string", "enum": ["low", "medium", "high"] }
                    },
                    "required": ["title", "priority"]
                }
            },
            { "name": "archive_issue", "inputSchema": {
                "type": "object",
                "properties": { "id": { "type": "string" } },
                "required": ["id"]
            }},
            { "name": "list_issues", "inputSchema": { "type": "object" } }
        ]})
    }

    /// Current catalog: `create_issue` swapped an enum value, `archive_issue`
    /// is gone, `add_comment` is new, `list_issues` is untouched.
    fn current() -> Value {
        json!({ "tools": [
            {
                "name": "create_issue",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "title": { "type": "string" },
                        "priority": { "type": "string", "enum": ["urgent", "medium", "high"] }
                    },
                    "required": ["title", "priority"]
                }
            },
            { "name": "add_comment", "inputSchema": {
                "type": "object",
                "properties": { "body": { "type": "string" } },
                "required": ["body"]
            }},
            { "name": "list_issues", "inputSchema": { "type": "object" } }
        ]})
    }

    /// A suite referencing the changed and removed tools plus a bystander
    /// test that must never be mentioned.
    fn suite() -> Value {
        json!({ "tools": [
            { "name": "create: hand authored", "tool": "create_issue" },
            { "name": "create: second hand authored", "tool": "create_issue" },
            { "name": "archive: hand authored", "tool": "archive_issue" },
            { "name": "list: untouched", "tool": "list_issues" }
        ]})
    }

    fn plan_with_suite() -> RegenPlan {
        let report = diff_tool_catalogs(&baseline(), &current()).unwrap();
        build_regen_plan(&report, &current(), Some(&suite()), "main")
    }

    #[test]
    fn changed_tool_lists_exactly_its_referencing_tests() {
        let plan = plan_with_suite();
        let item = plan
            .items
            .iter()
            .find(|i| i.kind == RegenKind::Changed)
            .expect("a changed item");
        assert_eq!(item.tool, "create_issue");
        assert_eq!(
            item.referencing_tests,
            vec![
                "create: hand authored".to_string(),
                "create: second hand authored".to_string()
            ]
        );
    }

    #[test]
    fn changed_tool_draft_regenerates_from_the_current_schema() {
        let plan = plan_with_suite();
        let item = plan
            .items
            .iter()
            .find(|i| i.kind == RegenKind::Changed)
            .expect("a changed item");
        let yaml = item.yaml.as_deref().expect("a draft");
        assert!(yaml.starts_with(&format!("  {REGEN_MARKER}\n")));
        // The synthesized enum value comes from the current catalog, not
        // the baseline: "low" no longer exists, "urgent" does.
        assert!(yaml.contains("priority: \"urgent\""));
        assert!(!yaml.contains("\"low\""));
    }

    #[test]
    fn removed_tool_gets_a_delete_or_rewrite_note_and_no_draft() {
        let plan = plan_with_suite();
        let item = plan
            .items
            .iter()
            .find(|i| i.kind == RegenKind::Removed)
            .expect("a removed item");
        assert_eq!(item.tool, "archive_issue");
        assert_eq!(
            item.referencing_tests,
            vec!["archive: hand authored".to_string()]
        );
        assert!(item.note.as_deref().unwrap().contains("delete or rewrite"));
        assert!(item.yaml.is_none());
        assert!(plan
            .combined_yaml()
            .contains("# removed tool `archive_issue`: delete or rewrite: archive: hand authored"));
    }

    #[test]
    fn added_tool_gets_a_fresh_draft_shaped_like_a_coverage_suggestion() {
        let plan = plan_with_suite();
        let item = plan
            .items
            .iter()
            .find(|i| i.kind == RegenKind::Added)
            .expect("an added item");
        assert_eq!(item.tool, "add_comment");
        let yaml = item.yaml.as_deref().expect("a draft");
        assert!(yaml.contains("- name: \"add_comment: valid arguments\""));
        assert!(yaml.contains("body: \"sample\""));
        assert!(yaml.contains("server: main"));
    }

    #[test]
    fn unaffected_tests_are_never_mentioned_anywhere_in_the_plan() {
        let plan = plan_with_suite();
        let serialized = serde_json::to_string(&plan).unwrap();
        assert!(!serialized.contains("list: untouched"));
        assert!(!plan.combined_yaml().contains("list: untouched"));
    }

    #[test]
    fn draft_names_dedup_against_existing_suite_names() {
        let suite = json!({ "tools": [
            { "name": "add_comment: valid arguments", "tool": "add_comment" }
        ]});
        let report = diff_tool_catalogs(&baseline(), &current()).unwrap();
        let plan = build_regen_plan(&report, &current(), Some(&suite), "main");
        let item = plan
            .items
            .iter()
            .find(|i| i.kind == RegenKind::Added)
            .expect("an added item");
        assert_eq!(
            item.draft_name.as_deref(),
            Some("add_comment: valid arguments (suggested)")
        );
    }

    #[test]
    fn without_a_suite_only_added_tool_drafts_are_emitted() {
        let report = diff_tool_catalogs(&baseline(), &current()).unwrap();
        let plan = build_regen_plan(&report, &current(), None, "server");
        assert!(plan.items.iter().all(|i| i.kind == RegenKind::Added));
        assert_eq!(plan.items.len(), 1);
        assert_eq!(plan.items[0].tool, "add_comment");
    }

    #[test]
    fn unsatisfiable_changed_schema_yields_a_note_instead_of_a_draft() {
        let old = json!({ "tools": [{ "name": "broken", "inputSchema": {
            "type": "object",
            "properties": { "n": { "type": "string" } },
            "required": ["n"]
        }}]});
        let new = json!({ "tools": [{ "name": "broken", "inputSchema": {
            "type": "object",
            "properties": { "n": { "type": "integer", "minimum": 10, "maximum": 3 } },
            "required": ["n"]
        }}]});
        let report = diff_tool_catalogs(&old, &new).unwrap();
        let suite = json!({ "tools": [{ "name": "broken: hand", "tool": "broken" }] });
        let plan = build_regen_plan(&report, &new, Some(&suite), "main");
        let item = &plan.items[0];
        assert_eq!(item.kind, RegenKind::Changed);
        assert!(item.yaml.is_none());
        assert!(item.note.as_deref().unwrap().contains("rewrite by hand"));
        assert_eq!(item.referencing_tests, vec!["broken: hand".to_string()]);
    }

    #[test]
    fn combined_yaml_parses_and_keeps_drafts_under_the_tools_key() {
        let plan = plan_with_suite();
        let combined = plan.combined_yaml();
        let parsed: serde_yaml::Value = serde_yaml::from_str(&combined).expect("combined parses");
        let tests = parsed
            .get("tools")
            .and_then(serde_yaml::Value::as_sequence)
            .expect("a tools sequence");
        // Two drafts (changed create_issue, added add_comment); the
        // removed tool contributes only a comment, never an entry.
        assert_eq!(tests.len(), 2);
        for test in tests {
            assert!(test.get("name").is_some());
            assert!(test.get("server").is_some());
            assert!(test.get("tool").is_some());
            assert!(test.get("expect").is_some());
        }
    }

    #[test]
    fn empty_report_yields_an_empty_plan_and_empty_yaml() {
        let report = diff_tool_catalogs(&baseline(), &baseline()).unwrap();
        let plan = build_regen_plan(&report, &baseline(), Some(&suite()), "main");
        assert!(plan.is_empty());
        assert_eq!(plan.combined_yaml(), "");
    }

    #[test]
    fn to_json_carries_combined_yaml_and_structured_items() {
        let plan = plan_with_suite();
        let wire = plan.to_json();
        assert_eq!(wire["items"][0]["kind"], "changed");
        assert_eq!(wire["items"][0]["tool"], "create_issue");
        assert!(wire["yaml"].as_str().unwrap().starts_with("tools:\n"));
    }

    use serde_json::json;
}
