mcp_typegen 0.1.0

use std::{fs, path::Path};

use schemars::schema::{Schema, SchemaObject};
use typify::{TypeSpace, TypeSpaceImpl, TypeSpaceSettings, UnknownPolicy};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let content = std::fs::read_to_string("./schema/schema.json").unwrap();
    let schema = serde_json::from_str::<schemars::schema::RootSchema>(&content).unwrap();

    // Create and configure settings
    let mut settings = TypeSpaceSettings::default();

    // Add a module prefix for all generated types
    settings.with_type_mod("mcp");

    // Add derives for all types
    settings
        .with_derive("Debug".to_string())
        .with_derive("Clone".to_string())
        .with_derive("PartialEq".to_string());

    // Enable struct builders
    //settings.with_struct_builder(true);

        // Create a schema object for the Result type
        let result_schema = SchemaObject {
            instance_type: Some(schemars::schema::InstanceType::Object.into()),
            object: Some(Box::new(schemars::schema::ObjectValidation {
                properties: {
                    let mut map = std::collections::BTreeMap::new();
                    map.insert("_meta".to_string(), Schema::Object(SchemaObject {
                        instance_type: Some(schemars::schema::InstanceType::Object.into()),
                        metadata: Some(Box::new(schemars::schema::Metadata {
                            description: Some("This result property is reserved by the protocol to allow clients and servers to attach additional metadata to their responses.".to_string()),
                            ..Default::default()
                        })),
                        ..Default::default()
                    }));
                    map
                },
                additional_properties: Some(Box::new(Schema::Object(SchemaObject::default()))),
                ..Default::default()
            })),
            ..Default::default()
        };

    // Handle special types that need replacement
    settings.with_replacement(
        "Result",
        "MCPResult",
        [TypeSpaceImpl::FromStr, TypeSpaceImpl::Display]
            .iter()
            .copied(),
    );

       // Add the Result schema conversion instead of replacement
       settings.with_conversion(
        result_schema,
        "MCPResult",
        [
            TypeSpaceImpl::FromStr,
            TypeSpaceImpl::Display,
            TypeSpaceImpl::Default,
        ].iter().copied(),
    );

    // Add any external crate dependencies
    settings.with_unknown_crates(UnknownPolicy::Allow);

    // Create type space with our settings
    let mut type_space = TypeSpace::new(&settings);

    // Add the schema
    type_space.add_root_schema(schema)?;

    // Generate the code
    let tokens = type_space.to_stream();
    let syntax_tree = syn::parse2::<syn::File>(tokens)?;
    let contents = prettyplease::unparse(&syntax_tree);

    // Add module-level attributes and documentation
    let final_contents = format!(
        r#"//! Model Context Protocol (MCP) Generated Types
//! 
//! This module contains automatically generated types from the MCP JSON schema.
//! Do not edit this file manually.

#![allow(clippy::all)]
#![allow(missing_docs)]
#![allow(dead_code)]

use serde::{{Serialize, Deserialize}};
use std::collections::HashMap;
use std::str::FromStr;

/// The base Result type for MCP operations.
/// This is renamed from 'Result' to avoid conflicts with the standard Result type.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MCPResult {{
    /// Optional metadata attached to the response
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub _meta: Option<HashMap<String, serde_json::Value>>,
    
    /// Additional properties that can be attached to the result
    #[serde(flatten)]
    pub extra: HashMap<String, serde_json::Value>,
}}

impl MCPResult {{
    pub fn new() -> Self {{
        Self {{
            _meta: None,
            extra: HashMap::new(),
        }}
    }}

    pub fn with_meta(meta: HashMap<String, serde_json::Value>) -> Self {{
        Self {{
            _meta: Some(meta),
            extra: HashMap::new(),
        }}
    }}
}}

    impl std::fmt::Display for MCPResult {{
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {{
            write!(f, "{{}}", serde_json::to_string(self).unwrap())
        }}
    }}

impl Default for MCPResult {{
    fn default() -> Self {{
        Self::new()
    }}
}}

    impl FromStr for MCPResult {{
        type Err = serde_json::Error;

        fn from_str(s: &str) -> Result<Self, Self::Err> {{
            serde_json::from_str(s)
        }}
    }}

/// Re-exports of commonly used types
pub mod prelude {{
    pub use super::{{
        ClientRequest, ServerRequest,
        ClientResult, ServerResult,
        ClientNotification, ServerNotification,
        SamplingMessage, PromptMessage,
        Resource, Tool, MCPResult,
    }};
}}

{contents}
"#
    );

    // Write the generated code
    let mut mod_gen = Path::new("src").to_path_buf();
    mod_gen.push("mcp_types");
    mod_gen.push("generated.rs");
    fs::write(&mod_gen, final_contents)?;

    // Generate mod.rs to handle public exports
    let mod_contents = r#"//! MCP types module
mod generated;
pub use generated::*;
pub use generated::prelude;
"#;

    let mut mod_file = Path::new("src").to_path_buf();
    mod_file.push("mcp_types");
    fs::create_dir_all(&mod_file)?;
    mod_file.push("mod.rs");
    fs::write(mod_file, mod_contents)?;

    // Ensure rebuild on schema changes
    println!("cargo:rerun-if-changed=schema/schema.json");

    Ok(())
}