llama-runner 2.3.2

use std::collections::BTreeMap;

use serde::Serialize;
use serde_json::Value;

use crate::mcp::{error::ParseToolError, gemma4::convert};

#[derive(Serialize, Clone)]
pub struct Gemma4Tool {
    pub function: Gemma4ToolFunction,
}

#[derive(Serialize, Clone, Default)]
pub struct Gemma4ToolFunction {
    pub name: String,
    pub description: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub parameters: Option<Gemma4ToolFunctionParams>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub response: Option<Gemma4ToolFunctionResponse>,
}

#[derive(Serialize, Clone, Default)]
pub struct Gemma4ToolFunctionParams {
    pub properties: BTreeMap<String, Gemma4ToolFunctionParamsProp>,
    pub required: Vec<String>,
}

#[derive(Serialize, Clone, Default)]
pub struct Gemma4ToolFunctionParamsProp {
    pub description: String,
    #[serde(rename = "type")]
    pub type_: Gemma4ToolFunctionParamsPropType,
    pub nullable: bool,
    /// String type only
    #[serde(rename = "enum", skip_serializing_if = "Vec::is_empty")]
    pub enum_: Vec<String>,
    /// Object type only
    #[serde(skip_serializing_if = "BTreeMap::is_empty")]
    pub properties: BTreeMap<String, Gemma4ToolFunctionParamsProp>,
    /// Object type only
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub required: Vec<String>,
    /// Array type only
    pub items: Option<Gemma4ToolFunctionParamsPropArrayItems>,
}

#[derive(Serialize, Clone, Default)]
pub struct Gemma4ToolFunctionParamsPropArrayItems {
    #[serde(skip_serializing_if = "BTreeMap::is_empty")]
    pub properties: BTreeMap<String, Gemma4ToolFunctionParamsProp>,
    #[serde(rename = "enum", skip_serializing_if = "Vec::is_empty")]
    pub required: Vec<String>,
    #[serde(rename = "type")]
    pub type_: Gemma4ToolFunctionParamsPropType,
}

#[derive(Serialize, Clone, Default)]
pub enum Gemma4ToolFunctionParamsPropType {
    String,
    Array,
    #[default]
    Object,
    Number,
    Boolean,
}

#[derive(Serialize, Clone, Default)]
pub struct Gemma4ToolFunctionResponse {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    #[serde(rename = "type")]
    pub type_: Option<Gemma4ToolFunctionParamsPropType>,
}

/// Following is generated by Claude
impl TryFrom<&rmcp::model::Tool> for Gemma4Tool {
    type Error = ParseToolError;

    fn try_from(value: &rmcp::model::Tool) -> Result<Self, Self::Error> {
        // The root schema object – needed for $ref resolution everywhere.
        let root = &value.input_schema;

        // ── required ──────────────────────────────────────────────────────────
        let required = if let Some(schema) = root.get("required").as_ref() {
            let members = schema
                .as_array()
                .ok_or(ParseToolError("input_schema.required".into()))?;

            members
                .iter()
                .enumerate()
                .map(|(idx, m)| {
                    m.as_str()
                        .ok_or(ParseToolError(
                            format!("input_schema.required[{idx}]").into(),
                        ))
                        .map(|s| s.to_string())
                })
                .collect::<Result<Vec<_>, _>>()?
        } else {
            vec![]
        };

        // ── properties ────────────────────────────────────────────────────────
        let properties = {
            let root = Value::Object(root.as_ref().clone());
            convert::convert_properties(&root, &root, "input_schema")?
        };

        let parameters = if root.is_empty() {
            None
        } else {
            Some(Gemma4ToolFunctionParams {
                properties,
                required,
            })
        };

        // ── response (output_schema) ──────────────────────────────────────────
        let response = value.output_schema.as_deref().map(|out| {
            let out_val = Value::Object(out.clone());

            let description = out_val
                .get("description")
                .and_then(|d| d.as_str())
                .map(str::to_string);

            // Determine the top-level type of the output schema.
            // Handles anyOf/oneOf nullable wrappers the same way as properties.
            let type_ = if let Some(variants) = out_val
                .get("anyOf")
                .or_else(|| out_val.get("oneOf"))
                .and_then(|v| v.as_array())
            {
                let (_nullable, effective) = convert::unwrap_any_of(variants, &out_val);
                effective
                    .and_then(|e| e.get("type"))
                    .and_then(|t| t.as_str())
                    .map(convert::map_type)
            } else {
                out_val
                    .get("type")
                    .and_then(|t| t.as_str())
                    .map(convert::map_type)
            };

            Gemma4ToolFunctionResponse { description, type_ }
        });

        Ok(Self {
            function: Gemma4ToolFunction {
                name: value.name.to_string(),
                description: value.description.as_deref().unwrap_or_default().to_string(),
                parameters,
                response,
            },
        })
    }
}

impl TryFrom<rmcp::model::Tool> for Gemma4Tool {
    type Error = ParseToolError;

    fn try_from(value: rmcp::model::Tool) -> Result<Self, Self::Error> {
        Self::try_from(&value)
    }
}

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

    fn make_tool(input_schema: Value, output_schema: Option<Value>) -> Tool {
        let mut tool = Tool::new(
            "my_tool",
            "does stuff",
            input_schema.as_object().unwrap().clone(),
        );
        tool.output_schema =
            output_schema.map(|v| std::sync::Arc::new(v.as_object().unwrap().clone()));
        tool
    }

    // ── Helper: run conversion and serialise to pretty JSON ──────────────────
    fn convert(input: Value, output: Option<Value>) -> serde_json::Value {
        let tool = make_tool(input, output);
        let gemma: Gemma4Tool = tool.try_into().expect("conversion failed");
        serde_json::to_value(gemma).expect("serialisation failed")
    }

    #[test]
    fn test_primitive_fields() {
        let schema = json!({
            "type": "object",
            "properties": {
                "my_bool": { "type": "boolean" },
                "my_int":  { "type": "integer", "format": "int32" }
            },
            "required": ["my_int", "my_bool"]
        });

        let v = convert(schema, None);
        let props = &v["function"]["parameters"]["properties"];
        assert_eq!(props["my_bool"]["type"], "Boolean");
        assert_eq!(props["my_int"]["type"], "Number");

        let req: Vec<&str> = v["function"]["parameters"]["required"]
            .as_array()
            .unwrap()
            .iter()
            .map(|v| v.as_str().unwrap())
            .collect();
        assert!(req.contains(&"my_int"));
        assert!(req.contains(&"my_bool"));
    }

    #[test]
    fn test_nullable_ref() {
        // my_nullable_enum: anyOf [ $ref, {type:null} ]
        let schema = json!({
            "type": "object",
            "properties": {
                "my_nullable_enum": {
                    "anyOf": [
                        { "$ref": "#/$defs/MyEnum" },
                        { "type": "null" }
                    ]
                }
            },
            "$defs": {
                "MyEnum": {
                    "oneOf": [
                        {
                            "type": "object",
                            "properties": { "StringNewType": { "type": "string" } },
                            "required": ["StringNewType"]
                        }
                    ]
                }
            }
        });

        let v = convert(schema, None);
        let prop = &v["function"]["parameters"]["properties"]["my_nullable_enum"];
        assert_eq!(prop["nullable"], true);
        // resolved through $ref + oneOf → first variant is an object
        assert_eq!(prop["type"], "Object");
    }

    #[test]
    fn test_array_of_numbers() {
        let schema = json!({
            "type": "object",
            "properties": {
                "floats": {
                    "type": "array",
                    "items": { "type": "number", "format": "float" }
                }
            },
            "required": ["floats"]
        });

        let v = convert(schema, None);
        let prop = &v["function"]["parameters"]["properties"]["floats"];
        assert_eq!(prop["type"], "Array");
        assert_eq!(prop["items"]["type"], "Number");
    }

    #[test]
    fn test_string_enum() {
        let schema = json!({
            "type": "object",
            "properties": {
                "color": {
                    "type": "string",
                    "enum": ["red", "green", "blue"]
                }
            }
        });

        let v = convert(schema, None);
        let prop = &v["function"]["parameters"]["properties"]["color"];
        assert_eq!(prop["type"], "String");
        let enums: Vec<&str> = prop["enum"]
            .as_array()
            .unwrap()
            .iter()
            .map(|v| v.as_str().unwrap())
            .collect();
        assert_eq!(enums, vec!["red", "green", "blue"]);
    }

    #[test]
    fn test_response_schema() {
        let input = json!({ "type": "object", "properties": {} });
        let output = json!({ "type": "string", "description": "the result" });

        let v = convert(input, Some(output));
        let resp = &v["function"]["response"];
        assert_eq!(resp["type"], "String");
        assert_eq!(resp["description"], "the result");
    }

    #[test]
    fn test_empty_schema_gives_no_parameters() {
        let v = convert(json!({}), None);
        assert!(v["function"]["parameters"].is_null());
    }

    #[test]
    fn test_full_sample_schema() {
        // The complete sample from the task description.
        let schema = json!({
            "$schema": "https://json-schema.org/draft/2020-12/schema",
            "title": "MyStruct",
            "type": "object",
            "properties": {
                "my_bool": { "type": "boolean" },
                "my_int":  { "type": "integer", "format": "int32" },
                "my_nullable_enum": {
                    "anyOf": [
                        { "$ref": "#/$defs/MyEnum" },
                        { "type": "null" }
                    ]
                }
            },
            "required": ["my_int", "my_bool"],
            "$defs": {
                "MyEnum": {
                    "oneOf": [
                        {
                            "type": "object",
                            "properties": {
                                "StringNewType": { "type": "string" }
                            },
                            "additionalProperties": false,
                            "required": ["StringNewType"]
                        },
                        {
                            "type": "object",
                            "properties": {
                                "StructVariant": {
                                    "type": "object",
                                    "properties": {
                                        "floats": {
                                            "type": "array",
                                            "items": { "type": "number", "format": "float" }
                                        }
                                    },
                                    "required": ["floats"]
                                }
                            },
                            "additionalProperties": false,
                            "required": ["StructVariant"]
                        }
                    ]
                }
            }
        });

        let v = convert(schema, None);
        let props = &v["function"]["parameters"]["properties"];

        assert_eq!(props["my_bool"]["type"], "Boolean");
        assert_eq!(props["my_int"]["type"], "Number");
        assert_eq!(props["my_nullable_enum"]["nullable"], true);
        // Resolved through $ref → oneOf → first variant (Object with StringNewType)
        assert_eq!(props["my_nullable_enum"]["type"], "Object");
        let inner = &props["my_nullable_enum"]["properties"]["StringNewType"];
        assert_eq!(inner["type"], "String");
    }
}