bevy_brp_mcp 0.22.2

//! BRP (Bevy Remote Protocol) client with unified execution interface
//!
//! This module provides a streamlined interface for communicating with BRP servers.
//! The `BrpClient` offers exactly 3 execution methods:
//! - `execute<R>()`: Primary API with automatic format discovery for result types that support it
//! - `execute_raw()`: Low-level API for debugging and format discovery engine
//! - `execute_streaming()`: Specialized API for watch operations with streaming responses

use reqwest::Response;
use serde_json::Value;
use tracing::warn;

use super::constants::BRP_EXTRAS_PREFIX;
use super::constants::ERROR_PATTERNS;
use super::constants::FORMAT_ERROR_HELP_FIELD;
use super::constants::FORMAT_ERROR_HELP_MESSAGE;
use super::constants::FORMAT_ERROR_ORIGINAL_ERROR_FIELD;
use super::constants::FORMAT_ERROR_SUGGESTED_ACTION;
use super::constants::FORMAT_ERROR_SUGGESTED_ACTION_FIELD;
use super::constants::FORMAT_ERROR_TYPE_GUIDE_FIELD;
use super::constants::JSON_RPC_ERROR_METHOD_NOT_FOUND;
use super::http_client::BrpHttpClient;
use super::operation::Operation;
use super::response_handling::BrpClientCallJsonResponse;
use super::response_handling::BrpClientError;
use super::response_handling::BrpToolConfig;
use super::response_handling::FormatCorrectionStatus;
use super::response_handling::ResponseStatus;
use super::response_handling::ResultStructBrpExt;
use crate::brp_tools::Port;
use crate::brp_tools::brp_type_guide;
use crate::error::Error;
use crate::error::Result;
use crate::tool::BrpMethod;
use crate::tool::ParameterName;

/// Client for executing a BRP operation
pub struct BrpClient {
    brp_method: BrpMethodName,
    port:       Port,
    params:     Option<Value>,
}

enum BrpMethodName {
    Known(BrpMethod),
    Application(String),
}

impl BrpMethodName {
    const fn as_str(&self) -> &str {
        match self {
            Self::Known(method) => method.as_str(),
            Self::Application(method) => method.as_str(),
        }
    }

    const fn known(&self) -> Option<BrpMethod> {
        match self {
            Self::Known(method) => Some(*method),
            Self::Application(_) => None,
        }
    }
}

impl BrpClient {
    /// Create a new BRP client for the given method, port, and parameters
    pub const fn new(brp_method: BrpMethod, port: Port, params: Option<Value>) -> Self {
        Self {
            brp_method: BrpMethodName::Known(brp_method),
            port,
            params,
        }
    }

    /// Create a BRP client for an application-defined method.
    pub const fn for_application(brp_method: String, port: Port, params: Option<Value>) -> Self {
        Self {
            brp_method: BrpMethodName::Application(brp_method),
            port,
            params,
        }
    }

    /// Primary execution method with automatic format discovery support
    ///
    /// This method implements the "execute-fail-discover" pattern:
    /// 1. Always executes the BRP request directly first
    /// 2. On success, returns the typed result immediately
    /// 3. On format errors, attempts format discovery if the result type supports it
    /// 4. Retries with corrected format if discovery succeeds
    ///
    /// Appending the type guide to an error is only attempted for result types with
    /// `BrpToolConfig::ADD_TYPE_GUIDE_TO_ERROR = true`. Result types with `ADD_TYPE_GUIDE_TO_ERROR
    /// = false` will return errors immediately without added `TypeGuide` .
    pub async fn execute<R>(&self) -> Result<R>
    where
        R: ResultStructBrpExt<
                Args = (
                    Option<Value>,
                    Option<Vec<Value>>,
                    Option<FormatCorrectionStatus>,
                ),
            > + BrpToolConfig
            + Send
            + 'static,
    {
        // ALWAYS execute direct first
        let direct_result = self.execute_direct_internal().await?;

        match direct_result {
            ResponseStatus::Success(data) => {
                // Success - no format discovery needed
                R::from_brp_client_response((
                    data,
                    None,
                    Some(FormatCorrectionStatus::NotAttempted),
                ))
            },
            ResponseStatus::Error(err) => {
                // Check if this result type supports adding the `TypeGuide`
                if R::ADD_TYPE_GUIDE_TO_ERROR && err.has_format_error_code() {
                    // embed type_guide information
                    self.try_add_type_guide_to_error(&err)
                        .await
                        .map_or_else(Err, |_| {
                            Err(Error::InvalidState(
                                "try_add_type_guide_to_error unexpectedly returned Ok".to_string(),
                            )
                            .into())
                        })
                } else {
                    // Regular error - enhance with context if possible
                    let enhanced_message =
                        self.enhance_error_message(err.get_message(), err.get_code());
                    Err(Error::tool_call_failed(enhanced_message).into())
                }
            },
        }
    }

    /// Low-level BRP execution without format discovery or result transformation
    ///
    /// This method provides direct access to BRP communication without any automatic
    /// format discovery or result type conversion. It returns raw `BrpClientResult`
    /// which can be either `Success(Option<Value>)` or `Error(BrpClientError)`.
    ///
    /// Primary use cases:
    /// - Debugging tools that need raw BRP responses (`brp_execute`)
    /// - Format discovery engine internal operations
    /// - Testing and diagnostic scenarios
    pub async fn execute_raw(&self) -> Result<ResponseStatus> {
        self.execute_direct_internal().await
    }

    /// Raw BRP execution without any error enhancement (used internally to prevent recursion)
    ///
    /// This method is identical to `execute_direct_internal()` but bypasses all error enhancement
    /// to prevent recursion when `TypeSchemaEngine` needs to fetch registry data.
    pub async fn execute_direct_internal_no_enhancement(&self) -> Result<ResponseStatus> {
        // Create HTTP client with our data
        let brp_http_client =
            BrpHttpClient::new(self.brp_method.as_str(), self.port, self.params.clone());

        // Send HTTP request (includes status check)
        let response = brp_http_client.send_request().await?;

        // Parse JSON-RPC response
        let brp_response = self.parse_json_response(response).await?;

        // Convert to BrpClientResult with special handling for bevy_brp_extras
        // NO ERROR ENHANCEMENT - return directly
        Ok(self.to_response_status(brp_response))
    }

    /// Execute the BRP request and return a streaming response
    ///
    /// This method is designed for watch operations that need to handle
    /// Server-Sent Events (SSE) streams. Unlike `execute()`, it:
    /// - Uses no timeout (streaming connections stay open)
    /// - Returns the raw response for the caller to process
    /// - Provides the same rich error context as other `BrpClient` methods
    pub async fn execute_streaming(&self) -> Result<Response> {
        // Create HTTP client with our data
        let brp_http_client =
            BrpHttpClient::new(self.brp_method.as_str(), self.port, self.params.clone());

        // Send HTTP request using streaming version (no timeout, includes status check)
        let response = brp_http_client.send_streaming_request().await?;

        Ok(response)
    }

    /// Internal direct execution - does the actual http call - we wanted the internal version so we
    /// can distinguish a canned call generated for a `ToolFn` by our macro, and the `execute_raw()`
    /// version we still allow to be called by bespoke tools like `brp_shutdown` and `brp_status`
    /// and the like.
    async fn execute_direct_internal(&self) -> Result<ResponseStatus> {
        // Create HTTP client with our data
        let brp_http_client =
            BrpHttpClient::new(self.brp_method.as_str(), self.port, self.params.clone());

        // Send HTTP request (includes status check)
        let response = brp_http_client.send_request().await?;

        // Parse JSON-RPC response
        let brp_response = self.parse_json_response(response).await?;

        // Convert to BrpClientResult with special handling for bevy_brp_extras
        Ok(self.to_response_status(brp_response))
    }

    /// Parse the JSON response from the BRP call to a running bevy app
    async fn parse_json_response(&self, response: Response) -> Result<BrpClientCallJsonResponse> {
        match response.json().await {
            Ok(json_response) => Ok(json_response),
            Err(e) => {
                warn!("BRP execute_brp_method: JSON parsing failed - error={e}");
                Err(
                    error_stack::Report::new(Error::JsonRpc("JSON parsing failed".to_string()))
                        .attach("Failed to parse BRP response JSON")
                        .attach(format!(
                            "Method: {}, Port: {}",
                            self.brp_method.as_str(),
                            self.port
                        ))
                        .attach(format!("Error: {e}")),
                )
            },
        }
    }

    /// Extract type names from BRP error messages using regex patterns
    fn extract_types_from_error_message(error_message: &str) -> Vec<String> {
        ERROR_PATTERNS
            .iter()
            .filter_map(|pattern| {
                regex::Regex::new(pattern)
                    .ok()
                    .and_then(|regex| regex.captures(error_message))
                    .and_then(|caps| caps.get(1))
                    .map(|m| (*m.as_str()).to_string())
            })
            .collect()
    }

    /// Enhance error messages with additional context when available
    ///
    /// Currently enhances:
    /// - `Entity` deserialization errors: Adds `Entity` IDs from parameters
    fn enhance_error_message(&self, original_message: &str, error_code: i32) -> String {
        // Check for entity deserialization errors
        if original_message.contains("Attempting to deserialize an invalid entity") {
            // Try to extract entity ID from parameters
            if let Some(params) = &self.params
                && let Some(entity_id) = params.get(ParameterName::Entity.as_ref())
            {
                return format!(
                    "Entity {entity_id} is not valid: {original_message} (error {error_code})"
                );
            }
        }

        // Default: return original message with error code
        format!("{original_message} (error {error_code})")
    }

    /// Enhanced format error creation with type guide embedding
    async fn try_add_type_guide_to_error(&self, error: &BrpClientError) -> Result<ResponseStatus> {
        // Step 1: Try parameter-based extraction using Operation enum
        let mut extracted_types = self
            .brp_method
            .known()
            .and_then(|method| Operation::try_from(method).ok())
            .map_or_else(Vec::new, |operation| {
                let params = self.params.as_ref().unwrap_or(&Value::Null);
                operation.extract_type_names(params)
            });

        // Step 2: Fallback to error message parsing if parameter extraction failed
        if extracted_types.is_empty() {
            extracted_types = Self::extract_types_from_error_message(error.get_message());
        }

        // Step 3: Handle results based on whether types were extracted
        if extracted_types.is_empty() {
            Self::create_minimal_type_error(error)
        } else {
            self.add_type_guide_to_error(error, extracted_types).await
        }
    }

    /// Create minimal error when no types can be extracted
    fn create_minimal_type_error(error: &BrpClientError) -> Result<ResponseStatus> {
        Err(Error::tool_call_failed_with_details(
            "Format error occurred but could not extract type information",
            serde_json::json!({
                FORMAT_ERROR_ORIGINAL_ERROR_FIELD: error.get_message(),
                FORMAT_ERROR_TYPE_GUIDE_FIELD: {
                    FORMAT_ERROR_HELP_FIELD: FORMAT_ERROR_HELP_MESSAGE,
                    FORMAT_ERROR_SUGGESTED_ACTION_FIELD: FORMAT_ERROR_SUGGESTED_ACTION
                }
            }),
        )
        .into())
    }

    /// Create full error with type guide embedded for extracted types
    async fn add_type_guide_to_error(
        &self,
        error: &BrpClientError,
        extracted_types: Vec<String>,
    ) -> Result<ResponseStatus> {
        let type_guide_response =
            brp_type_guide::generate_type_guide_response(self.port, &extracted_types).await?;

        Err(Error::tool_call_failed_with_details(
            "Format error - see 'type_guide' field for correct format",
            serde_json::json!({
                FORMAT_ERROR_ORIGINAL_ERROR_FIELD: error.get_message(),
                FORMAT_ERROR_TYPE_GUIDE_FIELD: type_guide_response
            }),
        )
        .into())
    }

    /// Convert the response JSON to a `ResponseStatus`
    fn to_response_status(&self, brp_response_json: BrpClientCallJsonResponse) -> ResponseStatus {
        if let Some(error) = brp_response_json.error {
            warn!(
                "BRP execute_brp_method: BRP returned error - code={}, message={}",
                error.code, error.message
            );

            // Check if this is a bevy_brp_extras method that's not found
            let enhanced_message = if error.code == JSON_RPC_ERROR_METHOD_NOT_FOUND {
                method_not_found_message(self.brp_method.as_str(), &error.message)
            } else {
                error.message
            };

            ResponseStatus::Error(BrpClientError {
                code:    error.code,
                message: enhanced_message,
                data:    error.data,
            })
        } else {
            ResponseStatus::Success(brp_response_json.result)
        }
    }
}

pub(crate) fn method_not_found_message(method: &str, message: &str) -> String {
    if method.starts_with(BRP_EXTRAS_PREFIX) {
        format!(
            "{message}. This method requires the bevy_brp_extras crate to be added to your Bevy app with the BrpExtrasPlugin"
        )
    } else {
        message.to_string()
    }
}