bevy_brp_mcp 0.22.2

use std::fs;
use std::time::SystemTime;
use std::time::UNIX_EPOCH;

use error_stack::Report;
use error_stack::ResultExt;
use rmcp::model::CallToolRequestParams;
use rmcp::model::CallToolResult;
use schemars::JsonSchema;
use serde_json::Map;
use serde_json::Value;
use serde_json::json;

use super::ParamStruct;
use super::ResultStruct;
use super::ToolDef;
use super::ToolResult;
use super::constants::CHARS_PER_TOKEN;
use super::constants::FILEPATH_FIELD;
use super::constants::INSTRUCTIONS_FIELD;
use super::constants::LARGE_RESPONSE_FILENAME_REPLACEMENT;
use super::constants::LARGE_RESPONSE_FILENAME_SANITIZE_CHARS;
use super::constants::LARGE_RESPONSE_INSTRUCTIONS;
use super::constants::ORIGINAL_SIZE_TOKENS_FIELD;
use super::constants::SAVED_TO_FILE_FIELD;
use super::json_response::AnySchemaValue;
use super::json_response::ToolCallJsonResponse;
use super::large_response::LargeResponseConfig;
use super::parameters;
use super::response_builder::Response;
use crate::error::Error;
use crate::error::Result;

/// Context passed to all handlers containing service, request, and MCP context
#[derive(Clone)]
pub struct HandlerContext {
    pub(super) tool_def: ToolDef,
    request:             CallToolRequestParams,
}

impl HandlerContext {
    /// Create a new `HandlerContext`
    pub(super) const fn new(tool_def: ToolDef, request: CallToolRequestParams) -> Self {
        Self { tool_def, request }
    }

    /// Common parameter extraction methods (used by both BRP and local handlers)
    pub(super) fn extract_parameter_values<T>(&self) -> Result<T>
    where
        T: serde::de::DeserializeOwned + JsonSchema,
    {
        // Get request arguments as JSON Value
        // Special case: if T is unit type, use null instead of empty object
        let args_value = if std::any::type_name::<T>() == "()" {
            Value::Null
        } else {
            self.request.arguments.as_ref().map_or_else(
                || Value::Object(Map::new()),
                |arguments| {
                    let mut arguments = arguments.clone();
                    parameters::normalize_arguments_for::<T>(&mut arguments);
                    Value::Object(arguments)
                },
            )
        };

        serde_json::from_value(args_value).map_err(|e| {
            tracing::debug!("Serde deserialization error: {e}");

            // Extract simplified type name (last component after ::)
            let type_name = std::any::type_name::<T>()
                .rsplit("::")
                .next()
                .unwrap_or("parameters");

            // Create user-friendly error message with serde details
            let user_message = format!(
                "Invalid parameter format for '{type_name}': {e}. Check the parameter types and structure match the tool's requirements"
            );

            error_stack::Report::new(Error::ParameterExtraction(user_message))
                .attach("Parameter validation failed")
                .attach(format!("Full type path: {}", std::any::type_name::<T>()))
                .attach(format!("Serde error details: {e}"))
        })
    }

    /// Get a field value from the request arguments
    ///
    /// This returns the raw request value without schema-based normalization.
    pub(super) fn extract_optional_named_field(&self, field_name: &str) -> Option<&Value> {
        self.request.arguments.as_ref()?.get(field_name)
    }

    /// Format a tool result into a `CallToolResult`
    pub(super) fn format_result<T, P>(&self, tool_result: ToolResult<T, P>) -> CallToolResult
    where
        T: ResultStruct,
        P: ParamStruct,
    {
        let tool_name = self.tool_def.tool_name;
        let call_info = tool_name.get_call_info();

        match tool_result.result {
            Ok(data) => match Response::success(&data, tool_result.params, call_info.clone(), self)
            {
                Ok(response) => {
                    // Handle large response here with access to tool_name
                    match self.handle_large_response_if_needed(response) {
                        Ok(processed) => processed.to_call_tool_result(),
                        Err(e) => Response::error_message(
                            format!("Failed to process response: {}", e.current_context()),
                            call_info,
                        )
                        .to_call_tool_result(),
                    }
                },
                Err(report) => Response::error_message(
                    format!("Internal error: {}", report.current_context()),
                    call_info,
                )
                .to_call_tool_result(),
            },
            Err(report) => match report.current_context() {
                Error::Structured { result } => {
                    // Create error response from structured result
                    match Response::error(
                        result.as_ref(),
                        tool_result.params,
                        call_info.clone(),
                        self,
                    ) {
                        Ok(response) => response.to_call_tool_result(),
                        Err(e) => Response::error_message(
                            format!("Failed to create error response: {}", e.current_context()),
                            call_info,
                        )
                        .to_call_tool_result(),
                    }
                },
                Error::ToolCall { message, details } => {
                    // Create error response with the error message and details
                    Response::error_with_details(message, details.as_ref(), call_info)
                        .to_call_tool_result()
                },
                _ => Response::error_message(
                    format!("Internal error: {}", report.current_context()),
                    call_info,
                )
                .to_call_tool_result(),
            },
        }
    }

    /// Format framework errors
    pub(super) fn format_framework_error(&self, error: Report<Error>) -> CallToolResult {
        let tool_name = self.tool_def.tool_name;
        let call_info = tool_name.get_call_info();

        Response::error_message(
            format!("Framework error: {}", error.current_context()),
            call_info,
        )
        .to_call_tool_result()
    }

    /// Handle large responses if needed
    fn handle_large_response_if_needed(
        &self,
        response: ToolCallJsonResponse,
    ) -> Result<ToolCallJsonResponse> {
        let large_response_config = LargeResponseConfig::default();

        // Check size and handle
        let response_json = serde_json::to_string(&response)
            .change_context(Error::General("Failed to serialize response".to_string()))?;
        let estimated_tokens = response_json.len() / CHARS_PER_TOKEN;

        if estimated_tokens > large_response_config.max_tokens
            && let Some(result_field) = &response.result
        {
            // Generate filename using self.tool_def.tool_name
            let timestamp = SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .change_context(Error::General("Failed to get timestamp".to_string()))?
                .as_secs();

            let sanitized_identifier = self.tool_def.tool_name.to_string().replace(
                LARGE_RESPONSE_FILENAME_SANITIZE_CHARS,
                LARGE_RESPONSE_FILENAME_REPLACEMENT,
            );
            let filename = format!(
                "{}{}{}.json",
                large_response_config.file_prefix, sanitized_identifier, timestamp
            );

            let filepath = large_response_config.temp_dir.join(&filename);

            let result_json = serde_json::to_string_pretty(result_field).change_context(
                Error::General("Failed to serialize result field".to_string()),
            )?;

            fs::write(&filepath, &result_json).change_context(Error::FileOperation(format!(
                "Failed to write result to {}",
                filepath.display()
            )))?;

            let mut modified_response = response;
            modified_response.result = Some(AnySchemaValue(json!({
                SAVED_TO_FILE_FIELD: true,
                FILEPATH_FIELD: filepath.to_string_lossy(),
                INSTRUCTIONS_FIELD: LARGE_RESPONSE_INSTRUCTIONS,
                ORIGINAL_SIZE_TOKENS_FIELD: estimated_tokens
            })));

            return Ok(modified_response);
        }

        Ok(response)
    }
}