use std::collections::BTreeMap;

use serde_json::{Value, json};
use sha2::{Digest, Sha256};

use super::{ToolDefinition, published_tool_definitions};

pub const PUBLIC_TOOL_SCHEMA_HASH_SCHEMA_VERSION: &str = "mcp_public_tool_schema_hashes.v1";

pub fn public_tool_schema_hashes() -> BTreeMap<String, String> {
    published_tool_definitions()
        .into_iter()
        .map(|tool| {
            let payload = json!({
                "schema_version": PUBLIC_TOOL_SCHEMA_HASH_SCHEMA_VERSION,
                "tool": tool_schema_hash_payload(&tool),
            });
            (
                tool.name.to_string(),
                sha256_prefixed(canonical_json(&payload).as_bytes()),
            )
        })
        .collect()
}

pub fn public_tool_schema_bundle_hash() -> String {
    let mut tools: Vec<Value> = published_tool_definitions()
        .iter()
        .map(tool_schema_hash_payload)
        .collect();
    tools.sort_by(|left, right| {
        left.get("name")
            .and_then(Value::as_str)
            .unwrap_or("")
            .cmp(right.get("name").and_then(Value::as_str).unwrap_or(""))
    });
    let payload = json!({
        "schema_version": PUBLIC_TOOL_SCHEMA_HASH_SCHEMA_VERSION,
        "surface": "mcp.task_first_vnext",
        "tools": tools,
    });
    sha256_prefixed(canonical_json(&payload).as_bytes())
}

fn tool_schema_hash_payload(tool: &ToolDefinition) -> Value {
    json!({
        "name": tool.name,
        "description": tool.description,
        "input_schema": tool.input_schema,
    })
}

fn canonical_json(value: &Value) -> String {
    match value {
        Value::Null => "null".to_string(),
        Value::Bool(true) => "true".to_string(),
        Value::Bool(false) => "false".to_string(),
        Value::Number(_) | Value::String(_) => serde_json::to_string(value)
            .expect("JSON scalars generated by serde_json should serialize"),
        Value::Array(items) => {
            let inner = items
                .iter()
                .map(canonical_json)
                .collect::<Vec<_>>()
                .join(",");
            format!("[{inner}]")
        }
        Value::Object(map) => {
            let mut keys = map.keys().collect::<Vec<_>>();
            keys.sort();
            let inner = keys
                .into_iter()
                .map(|key| {
                    let key_json = serde_json::to_string(key)
                        .expect("JSON object keys should serialize as strings");
                    format!("{}:{}", key_json, canonical_json(&map[key]))
                })
                .collect::<Vec<_>>()
                .join(",");
            format!("{{{inner}}}")
        }
    }
}

fn sha256_prefixed(bytes: &[u8]) -> String {
    let mut hasher = Sha256::new();
    hasher.update(bytes);
    format!("sha256:{}", hex::encode(hasher.finalize()))
}
