//! Configuration object accepted by the `codex` MCP tool-call.

use codex_arg0::Arg0DispatchPaths;
use codex_core::config::Config;
use codex_core::config::ConfigBuilder;
use codex_core::config::ConfigOverrides;
use codex_protocol::ThreadId;
use codex_protocol::config_types::SandboxMode;
use codex_protocol::protocol::AskForApproval;
use codex_utils_json_to_toml::json_to_toml;
use rmcp::model::JsonObject;
use rmcp::model::Tool;
use schemars::JsonSchema;
use schemars::r#gen::SchemaSettings;
use serde::Deserialize;
use serde::Serialize;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;

/// Client-supplied configuration for a `codex` tool-call.
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, Default)]
#[serde(rename_all = "kebab-case", deny_unknown_fields)]
#[schemars(deny_unknown_fields)]
pub struct CodexToolCallParam {
    /// The *initial user prompt* to start the Codex conversation.
    pub prompt: String,

    /// Optional override for the model name (e.g. 'gpt-5.2', 'gpt-5.2-codex').
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub model: Option<String>,

    /// Working directory for the session. If relative, it is resolved against
    /// the server process's current working directory.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cwd: Option<String>,

    /// Approval policy for shell commands generated by the model:
    /// `untrusted`, `on-failure`, `on-request`, `never`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub approval_policy: Option<CodexToolCallApprovalPolicy>,

    /// Sandbox mode: `read-only`, `workspace-write`, or `danger-full-access`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub sandbox: Option<CodexToolCallSandboxMode>,

    /// Individual config settings that will override what is in
    /// CODEX_HOME/config.toml.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub config: Option<HashMap<String, serde_json::Value>>,

    /// The set of instructions to use instead of the default ones.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub base_instructions: Option<String>,

    /// Developer instructions that should be injected as a developer role message.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub developer_instructions: Option<String>,

    /// Prompt used when compacting the conversation.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub compact_prompt: Option<String>,
}

/// Custom enum mirroring [`AskForApproval`], but has an extra dependency on
/// [`JsonSchema`].
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
#[serde(rename_all = "kebab-case")]
pub enum CodexToolCallApprovalPolicy {
    Untrusted,
    OnFailure,
    OnRequest,
    Never,
}

impl From<CodexToolCallApprovalPolicy> for AskForApproval {
    fn from(value: CodexToolCallApprovalPolicy) -> Self {
        match value {
            CodexToolCallApprovalPolicy::Untrusted => AskForApproval::UnlessTrusted,
            CodexToolCallApprovalPolicy::OnFailure => AskForApproval::OnFailure,
            CodexToolCallApprovalPolicy::OnRequest => AskForApproval::OnRequest,
            CodexToolCallApprovalPolicy::Never => AskForApproval::Never,
        }
    }
}

/// Custom enum mirroring [`SandboxMode`] from config_types.rs, but with
/// `JsonSchema` support.
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
#[serde(rename_all = "kebab-case")]
pub enum CodexToolCallSandboxMode {
    ReadOnly,
    WorkspaceWrite,
    DangerFullAccess,
}

impl From<CodexToolCallSandboxMode> for SandboxMode {
    fn from(value: CodexToolCallSandboxMode) -> Self {
        match value {
            CodexToolCallSandboxMode::ReadOnly => SandboxMode::ReadOnly,
            CodexToolCallSandboxMode::WorkspaceWrite => SandboxMode::WorkspaceWrite,
            CodexToolCallSandboxMode::DangerFullAccess => SandboxMode::DangerFullAccess,
        }
    }
}

/// Builds a `Tool` definition (JSON schema etc.) for the Codex tool-call.
pub(crate) fn create_tool_for_codex_tool_call_param() -> Tool {
    let schema = SchemaSettings::draft2019_09()
        .with(|s| {
            s.inline_subschemas = true;
            s.option_add_null_type = false;
        })
        .into_generator()
        .into_root_schema_for::<CodexToolCallParam>();

    let input_schema = create_tool_input_schema(schema, "Codex tool schema should serialize");

    Tool::new(
        "codex",
        "Run a Codex session. Accepts configuration parameters matching the Codex Config struct.",
        input_schema,
    )
    .with_title("Codex")
    .with_raw_output_schema(codex_tool_output_schema())
}

fn codex_tool_output_schema() -> Arc<JsonObject> {
    let schema = serde_json::json!({
        "type": "object",
        "properties": {
            "threadId": { "type": "string" },
            "content": { "type": "string" }
        },
        "required": ["threadId", "content"],
    });
    match schema {
        serde_json::Value::Object(map) => Arc::new(map),
        _ => unreachable!("json literal must be an object"),
    }
}

impl CodexToolCallParam {
    /// Returns the initial user prompt to start the Codex conversation and the
    /// effective Config object generated from the supplied parameters.
    pub async fn into_config(
        self,
        arg0_paths: Arg0DispatchPaths,
    ) -> std::io::Result<(String, Config)> {
        let Self {
            prompt,
            model,
            cwd,
            approval_policy,
            sandbox,
            config: cli_overrides,
            base_instructions,
            developer_instructions,
            compact_prompt,
        } = self;

        // Build the `ConfigOverrides` recognized by codex-core.
        let overrides = ConfigOverrides {
            model,
            cwd: cwd.map(PathBuf::from),
            approval_policy: approval_policy.map(Into::into),
            sandbox_mode: sandbox.map(Into::into),
            codex_self_exe: arg0_paths.codex_self_exe.clone(),
            codex_linux_sandbox_exe: arg0_paths.codex_linux_sandbox_exe.clone(),
            main_execve_wrapper_exe: arg0_paths.main_execve_wrapper_exe.clone(),
            base_instructions,
            developer_instructions,
            compact_prompt,
            ..Default::default()
        };

        let cli_overrides = cli_overrides
            .unwrap_or_default()
            .into_iter()
            .map(|(k, v)| (k, json_to_toml(v)))
            .collect();

        let cfg = ConfigBuilder::default()
            .cli_overrides(cli_overrides)
            .harness_overrides(overrides)
            .build()
            .await?;

        Ok((prompt, cfg))
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct CodexToolCallReplyParam {
    /// DEPRECATED: use threadId instead.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    conversation_id: Option<String>,

    /// The thread id for this Codex session.
    /// This field is required, but we keep it optional here for backward
    /// compatibility for clients that still use conversationId.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    thread_id: Option<String>,

    /// The *next user prompt* to continue the Codex conversation.
    pub prompt: String,
}

impl CodexToolCallReplyParam {
    pub(crate) fn get_thread_id(&self) -> anyhow::Result<ThreadId> {
        if let Some(thread_id) = &self.thread_id {
            let thread_id = ThreadId::from_string(thread_id)?;
            Ok(thread_id)
        } else if let Some(conversation_id) = &self.conversation_id {
            let thread_id = ThreadId::from_string(conversation_id)?;
            Ok(thread_id)
        } else {
            Err(anyhow::anyhow!(
                "either threadId or conversationId must be provided"
            ))
        }
    }
}

/// Builds a `Tool` definition for the `codex-reply` tool-call.
pub(crate) fn create_tool_for_codex_tool_call_reply_param() -> Tool {
    let schema = SchemaSettings::draft2019_09()
        .with(|s| {
            s.inline_subschemas = true;
            s.option_add_null_type = false;
        })
        .into_generator()
        .into_root_schema_for::<CodexToolCallReplyParam>();

    let input_schema = create_tool_input_schema(schema, "Codex reply tool schema should serialize");

    Tool::new(
        "codex-reply",
        "Continue a Codex conversation by providing the thread id and prompt.",
        input_schema,
    )
    .with_title("Codex Reply")
    .with_raw_output_schema(codex_tool_output_schema())
}

fn create_tool_input_schema(
    schema: schemars::schema::RootSchema,
    panic_message: &str,
) -> Arc<JsonObject> {
    #[expect(clippy::expect_used)]
    let schema_value = serde_json::to_value(&schema).expect(panic_message);
    let mut schema_object = match schema_value {
        serde_json::Value::Object(object) => object,
        _ => panic!("tool schema should serialize to a JSON object"),
    };

    // Prefer keeping the "core" JSON Schema keys while still preserving `$defs`
    // in case any `$ref` leaks into the generated schema (even though we try
    // to inline subschemas).
    let mut input_schema = JsonObject::new();
    for key in [
        "additionalProperties",
        "properties",
        "required",
        "type",
        "$defs",
        "definitions",
    ] {
        if let Some(value) = schema_object.remove(key) {
            input_schema.insert(key.to_string(), value);
        }
    }

    Arc::new(input_schema)
}

#[cfg(test)]
mod tests {
    use super::*;
    use pretty_assertions::assert_eq;

    /// We include a test to verify the exact JSON schema as "executable
    /// documentation" for the schema. When can track changes to this test as a
    /// way to audit changes to the generated schema.
    ///
    /// Seeing the fully expanded schema makes it easier to casually verify that
    /// the generated JSON for enum types such as "approval-policy" is compact.
    /// Ideally, modelcontextprotocol/inspector would provide a simpler UI for
    /// enum fields versus open string fields to take advantage of this.
    ///
    /// As of 2025-05-04, there is an open PR for this:
    /// https://github.com/modelcontextprotocol/inspector/pull/196
    #[test]
    fn verify_codex_tool_json_schema() {
        let tool = create_tool_for_codex_tool_call_param();
        let tool_json = serde_json::to_value(&tool).expect("tool serializes");
        let expected_tool_json = serde_json::json!({
          "description": "Run a Codex session. Accepts configuration parameters matching the Codex Config struct.",
          "inputSchema": {
            "additionalProperties": false,
            "properties": {
              "approval-policy": {
                "description": "Approval policy for shell commands generated by the model: `untrusted`, `on-failure`, `on-request`, `never`.",
                "enum": [
                  "untrusted",
                  "on-failure",
                  "on-request",
                  "never"
                ],
                "type": "string"
              },
              "base-instructions": {
                "description": "The set of instructions to use instead of the default ones.",
                "type": "string"
              },
              "compact-prompt": {
                "description": "Prompt used when compacting the conversation.",
                "type": "string"
              },
              "config": {
                "additionalProperties": true,
                "description": "Individual config settings that will override what is in CODEX_HOME/config.toml.",
                "type": "object"
              },
              "cwd": {
                "description": "Working directory for the session. If relative, it is resolved against the server process's current working directory.",
                "type": "string"
              },
              "developer-instructions": {
                "description": "Developer instructions that should be injected as a developer role message.",
                "type": "string"
              },
              "model": {
                "description": "Optional override for the model name (e.g. 'gpt-5.2', 'gpt-5.2-codex').",
                "type": "string"
              },
              "prompt": {
                "description": "The *initial user prompt* to start the Codex conversation.",
                "type": "string"
              },
              "sandbox": {
                "description": "Sandbox mode: `read-only`, `workspace-write`, or `danger-full-access`.",
                "enum": [
                  "read-only",
                  "workspace-write",
                  "danger-full-access"
                ],
                "type": "string"
              }
            },
            "required": [
              "prompt"
            ],
            "type": "object"
          },
          "name": "codex",
          "outputSchema": {
            "properties": {
              "content": {
                "type": "string"
              },
              "threadId": {
                "type": "string"
              }
            },
            "required": [
              "threadId",
              "content"
            ],
            "type": "object"
          },
          "title": "Codex"
        });
        assert_eq!(expected_tool_json, tool_json);
    }

    #[test]
    fn codex_tool_call_param_rejects_removed_profile_field() {
        let err = serde_json::from_value::<CodexToolCallParam>(serde_json::json!({
            "prompt": "hello",
            "profile": "work"
        }))
        .expect_err("removed profile field should fail");

        assert!(
            err.to_string().contains("unknown field `profile`"),
            "unexpected error: {err}"
        );
    }

    #[test]
    fn verify_codex_tool_reply_json_schema() {
        let tool = create_tool_for_codex_tool_call_reply_param();
        let tool_json = serde_json::to_value(&tool).expect("tool serializes");
        let expected_tool_json = serde_json::json!({
          "description": "Continue a Codex conversation by providing the thread id and prompt.",
          "inputSchema": {
            "properties": {
              "conversationId": {
                "description": "DEPRECATED: use threadId instead.",
                "type": "string"
              },
              "prompt": {
                "description": "The *next user prompt* to continue the Codex conversation.",
                "type": "string"
              },
              "threadId": {
                "description": "The thread id for this Codex session. This field is required, but we keep it optional here for backward compatibility for clients that still use conversationId.",
                "type": "string"
              }
            },
            "required": [
              "prompt",
            ],
            "type": "object",
          },
          "name": "codex-reply",
          "outputSchema": {
            "properties": {
              "content": {
                "type": "string"
              },
              "threadId": {
                "type": "string"
              }
            },
            "required": [
              "threadId",
              "content"
            ],
            "type": "object"
          },
          "title": "Codex Reply",
        });
        assert_eq!(expected_tool_json, tool_json);
    }
}
