use super::augment_tool_spec_for_code_mode as augment_tool_spec_with_budget;
use super::code_mode_name_for_tool_name;
use super::tool_spec_to_code_mode_tool_definition;
use crate::AdditionalProperties;
use crate::FreeformTool;
use crate::FreeformToolFormat;
use crate::JsonSchema;
use crate::ResponsesApiNamespace;
use crate::ResponsesApiNamespaceTool;
use crate::ResponsesApiTool;
use crate::ToolName;
use crate::ToolSpec;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::collections::BTreeMap;

fn augment_tool_spec_for_code_mode(spec: ToolSpec) -> ToolSpec {
    augment_tool_spec_with_budget(spec, /*code_mode_input_schema_max_bytes*/ None)
}

#[test]
fn code_mode_input_schema_uses_the_larger_budget() {
    let default_code_mode = codex_code_mode::DEFAULT_INPUT_SCHEMA_MAX_BYTES;
    for (mcp, code_mode, expected) in [
        (None, None, default_code_mode),
        (Some(5_000), Some(1_000), default_code_mode),
        (None, Some(1_000), default_code_mode),
        (Some(30_000), Some(20_000), 30_000),
        (Some(20_000), Some(30_000), 30_000),
    ] {
        assert_eq!(
            super::effective_input_schema_max_bytes(mcp, code_mode),
            expected
        );
    }
}

#[test]
fn code_mode_materializes_mcp_output_schemas() {
    let output = json!({"type": "object", "properties": {"ok": {"type": "boolean"}}});
    let mut tool = rmcp::model::Tool::new(
        "lookup_order",
        "Look up an order",
        std::sync::Arc::new(rmcp::model::object(json!({"type": "object"}))),
    );
    tool.output_schema = Some(std::sync::Arc::new(rmcp::model::object(output.clone())));
    let parsed = crate::mcp_tool_to_responses_api_tool(
        &ToolName::plain("lookup_order"),
        &tool,
        /*schema_max_bytes*/ None,
    )
    .unwrap();
    let eager = ResponsesApiTool {
        output_schema: Some(crate::mcp_call_tool_result_output_schema(output).into()),
        ..parsed.clone()
    };
    assert_eq!(
        super::collect_code_mode_tool_definitions(
            [&ToolSpec::Function(parsed)],
            /*code_mode_input_schema_max_bytes*/ None,
            /*tool_description_first*/ false,
        ),
        super::collect_code_mode_tool_definitions(
            [&ToolSpec::Function(eager)],
            /*code_mode_input_schema_max_bytes*/ None,
            /*tool_description_first*/ false,
        ),
    );
}

#[test]
fn code_mode_tool_names_do_not_prefix_the_default_namespace() {
    for tool_name in [
        ToolName::plain("apply_patch"),
        ToolName::namespaced("functions", "apply_patch"),
    ] {
        assert_eq!(code_mode_name_for_tool_name(&tool_name), "apply_patch");
    }

    assert_eq!(
        code_mode_name_for_tool_name(&ToolName::namespaced("editor", "apply_patch")),
        "editor__apply_patch"
    );
}

#[test]
fn augment_tool_spec_for_code_mode_augments_function_tools() {
    assert_eq!(
        augment_tool_spec_for_code_mode(ToolSpec::Function(ResponsesApiTool {
            name: "lookup_order".to_string(),
            description: "Look up an order".to_string(),
            strict: false,
            defer_loading: Some(true),
            parameters: JsonSchema::object(
                BTreeMap::from([(
                    "order_id".to_string(),
                    JsonSchema::string(/*description*/ None),
                )]),
                Some(vec!["order_id".to_string()]),
                Some(AdditionalProperties::Boolean(false))
            ),
            output_schema: Some(
                json!({
                    "type": "object",
                    "properties": {
                        "ok": {"type": "boolean"}
                    },
                    "required": ["ok"],
                })
                .into()
            ),
        })),
        ToolSpec::Function(ResponsesApiTool {
            name: "lookup_order".to_string(),
            description: r#"Look up an order

exec tool declaration:
```ts
declare const tools: { lookup_order(args: { order_id: string; }): Promise<{ ok: boolean; }>; };
```"#
                .to_string(),
            strict: false,
            defer_loading: Some(true),
            parameters: JsonSchema::object(
                BTreeMap::from([(
                    "order_id".to_string(),
                    JsonSchema::string(/*description*/ None),
                )]),
                Some(vec!["order_id".to_string()]),
                Some(AdditionalProperties::Boolean(false))
            ),
            output_schema: Some(
                json!({
                    "type": "object",
                    "properties": {
                        "ok": {"type": "boolean"}
                    },
                    "required": ["ok"],
                })
                .into()
            ),
        })
    );
}

#[test]
fn augment_tool_spec_for_code_mode_preserves_exec_tool_description() {
    assert_eq!(
        augment_tool_spec_for_code_mode(ToolSpec::Freeform(FreeformTool {
            name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(),
            description: "Run code".to_string(),
            defer_loading: None,
            format: FreeformToolFormat {
                r#type: "grammar".to_string(),
                syntax: "lark".to_string(),
                definition: "start: \"exec\"".to_string(),
            },
        })),
        ToolSpec::Freeform(FreeformTool {
            name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(),
            description: "Run code".to_string(),
            defer_loading: None,
            format: FreeformToolFormat {
                r#type: "grammar".to_string(),
                syntax: "lark".to_string(),
                definition: "start: \"exec\"".to_string(),
            },
        })
    );
}

#[test]
fn tool_spec_to_code_mode_tool_definition_returns_augmented_nested_tools() {
    let spec = ToolSpec::Freeform(FreeformTool {
        name: "apply_patch".to_string(),
        description: "Apply a patch".to_string(),
        defer_loading: None,
        format: FreeformToolFormat {
            r#type: "grammar".to_string(),
            syntax: "lark".to_string(),
            definition: "start: \"patch\"".to_string(),
        },
    });

    assert_eq!(
        tool_spec_to_code_mode_tool_definition(&spec),
        Some(codex_code_mode::ToolDefinition {
            name: "apply_patch".to_string(),
            tool_name: ToolName::plain("apply_patch"),
            description: r#"Apply a patch

exec tool declaration:
```ts
declare const tools: { apply_patch(input: string): Promise<unknown>; };
```"#
                .to_string(),
            kind: codex_code_mode::CodeModeToolKind::Freeform,
            input_schema: None,
            input_schema_max_bytes: None,
            output_schema: None,
        })
    );
}

#[test]
fn tool_spec_to_code_mode_tool_definition_supports_namespaced_custom_tools() {
    let spec = ToolSpec::Namespace(ResponsesApiNamespace {
        name: "editor".to_string(),
        description: "Editing tools".to_string(),
        tools: vec![ResponsesApiNamespaceTool::Custom(FreeformTool {
            name: "apply_patch".to_string(),
            description: "Apply a patch".to_string(),
            defer_loading: None,
            format: FreeformToolFormat {
                r#type: "grammar".to_string(),
                syntax: "lark".to_string(),
                definition: "start: \"patch\"".to_string(),
            },
        })],
    });

    assert_eq!(
        tool_spec_to_code_mode_tool_definition(&spec),
        Some(codex_code_mode::ToolDefinition {
            name: "editor__apply_patch".to_string(),
            tool_name: ToolName::namespaced("editor", "apply_patch"),
            description: r#"Apply a patch

exec tool declaration:
```ts
declare const tools: { editor__apply_patch(input: string): Promise<unknown>; };
```"#
                .to_string(),
            kind: codex_code_mode::CodeModeToolKind::Freeform,
            input_schema: None,
            input_schema_max_bytes: None,
            output_schema: None,
        })
    );
}

#[test]
fn tool_search_spec_converts_to_a_callable_code_mode_definition() {
    let definition = tool_spec_to_code_mode_tool_definition(&ToolSpec::ToolSearch {
        execution: "client".to_string(),
        description: "Native tool discovery".to_string(),
        parameters: JsonSchema::object(
            BTreeMap::from([(
                "query".to_string(),
                JsonSchema::string(/*description*/ None),
            )]),
            Some(vec!["query".to_string()]),
            Some(AdditionalProperties::Boolean(false)),
        ),
    })
    .expect("tool search should be callable from Code Mode");
    assert_eq!(
        (
            &definition.name,
            &definition.tool_name,
            &definition.input_schema,
            definition.input_schema_max_bytes,
        ),
        (
            &"tool_search".to_string(),
            &ToolName::plain("tool_search"),
            &Some(json!({
                "type": "object",
                "properties": {
                    "query": { "type": "string" },
                },
                "required": ["query"],
                "additionalProperties": false,
            })),
            Some(codex_code_mode::DEFAULT_INPUT_SCHEMA_MAX_BYTES),
        ),
    );
    assert!(
        definition
            .description
            .contains("tool_search(args: { query: string; })")
    );
    assert!(
        definition
            .description
            .contains("tools[result.name](arguments)")
    );
    assert!(!definition.description.contains("Native tool discovery"));
}