text-to-cypher 0.2.2

use crate::chat::{ChatMessage, ChatRequest, ChatRole};
use crate::mcp::tools::TextToCypherTool;
use crate::usage::TokenUsage;
use async_trait::async_trait;
use futures_util::StreamExt;
use rust_mcp_sdk::schema::TextContent;
use rust_mcp_sdk::schema::{
    CallToolRequest, CallToolResult, ListResourcesRequest, ListResourcesResult, ListToolsRequest, ListToolsResult,
    ReadResourceRequest, ReadResourceResult, Resource, RpcError, TextResourceContents, schema_utils::CallToolError,
};
use rust_mcp_sdk::{McpServer, mcp_server::ServerHandler};
use std::fmt::Write;
use std::sync::Arc;

// Custom Handler to handle MCP Messages
pub struct MyServerHandler;

#[async_trait]
impl ServerHandler for MyServerHandler {
    async fn handle_list_tools_request(
        &self,
        _request: ListToolsRequest,
        _runtime: Arc<dyn McpServer>,
    ) -> std::result::Result<ListToolsResult, RpcError> {
        tracing::info!("Handling List Tools Request");
        Ok(ListToolsResult {
            meta: None,
            next_cursor: None,
            tools: vec![TextToCypherTool::tool()],
        })
    }

    async fn handle_list_resources_request(
        &self,
        _request: ListResourcesRequest,
        _runtime: Arc<dyn McpServer>,
    ) -> std::result::Result<ListResourcesResult, RpcError> {
        tracing::info!("Handling List Resources Request");

        match get_falkordb_graphs().await {
            Ok(graphs) => {
                let resources: Vec<Resource> = graphs
                    .into_iter()
                    .map(|graph_name| Resource {
                        uri: format!("falkordb://graph/{graph_name}"),
                        name: format!("Graph: {graph_name}"),
                        description: Some(format!("FalkorDB graph database: {graph_name}")),
                        mime_type: Some("application/json".to_string()),
                        annotations: None,
                        meta: None,
                        size: None,
                        title: None,
                    })
                    .collect();

                Ok(ListResourcesResult {
                    meta: None,
                    next_cursor: None,
                    resources,
                })
            }
            Err(e) => {
                tracing::error!("Failed to list FalkorDB graphs: {}", e);
                Err(RpcError::internal_error())
            }
        }
    }

    async fn handle_read_resource_request(
        &self,
        request: ReadResourceRequest,
        _runtime: Arc<dyn McpServer>,
    ) -> std::result::Result<ReadResourceResult, RpcError> {
        tracing::info!("Handling Read Resource Request for URI: {}", request.params.uri);

        // Parse the URI to extract graph name
        if let Some(graph_name) = request.params.uri.strip_prefix("falkordb://graph/") {
            match get_graph_schema_via_api(graph_name).await {
                Ok(schema_info) => {
                    let text_content = TextResourceContents {
                        uri: request.params.uri,
                        mime_type: Some("application/json".to_string()),
                        text: schema_info,
                        meta: None,
                    };
                    Ok(ReadResourceResult {
                        meta: None,
                        contents: vec![text_content.into()],
                    })
                }
                Err(e) => {
                    tracing::error!("Failed to read graph schema for {}: {}", graph_name, e);
                    Err(RpcError::invalid_params())
                }
            }
        } else {
            Err(RpcError::invalid_params())
        }
    }

    async fn handle_call_tool_request(
        &self,
        request: CallToolRequest,
        _runtime: Arc<dyn McpServer>,
    ) -> std::result::Result<CallToolResult, CallToolError> {
        tracing::info!("Handling Call Tool Request");
        if request.tool_name() == TextToCypherTool::tool_name() {
            // Get the arguments from the request
            let arguments = request.params.arguments.unwrap_or_default();
            let arguments_value = serde_json::Value::Object(arguments);

            // Parse the tool arguments
            match serde_json::from_value::<TextToCypherTool>(arguments_value.clone()) {
                Ok(tool_args) => {
                    tracing::info!("TextToCypherTool called with arguments:");
                    tracing::info!("  graph_name: {}", tool_args.graph_name);
                    tracing::info!("  question: {}", tool_args.question);

                    // Forward the request to the HTTP endpoint
                    match forward_to_http_endpoint(tool_args).await {
                        Ok(result) => Ok(CallToolResult::text_content(vec![TextContent::from(result)])),
                        Err(e) => {
                            tracing::error!("Failed to forward request to HTTP endpoint: {}", e);
                            Err(CallToolError::new(std::io::Error::other(format!(
                                "HTTP forwarding failed: {e}"
                            ))))
                        }
                    }
                }
                Err(e) => {
                    tracing::error!("Failed to parse TextToCypherTool arguments: {}", e);
                    Err(CallToolError::new(e))
                }
            }
        } else {
            Err(CallToolError::unknown_tool(request.tool_name().to_string()))
        }
    }

    async fn on_initialized(
        &self,
        _runtime: Arc<dyn McpServer>,
    ) {
    }
}

// Helper function to forward MCP tool request to HTTP endpoint
async fn forward_to_http_endpoint(
    tool_args: TextToCypherTool
) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
    let http_request = create_http_request_payload(tool_args);
    let response = send_http_request(&http_request).await?;
    process_sse_response(response).await
}

// Create HTTP request payload for the text-to-cypher endpoint
fn create_http_request_payload(tool_args: TextToCypherTool) -> serde_json::Value {
    let chat_request = ChatRequest {
        messages: vec![ChatMessage {
            role: ChatRole::User,
            content: tool_args.question,
        }],
    };

    serde_json::json!({
        "graph_name": tool_args.graph_name,
        "chat_request": chat_request,
        "model": null,
        "key": null
    })
}

// Send HTTP request to the text-to-cypher endpoint
async fn send_http_request(
    http_request: &serde_json::Value
) -> Result<reqwest::Response, Box<dyn std::error::Error + Send + Sync>> {
    let client = reqwest::Client::new();
    let response = client
        .post("http://127.0.0.1:8080/text_to_cypher")
        .header("Content-Type", "application/json")
        .json(http_request)
        .send()
        .await?;

    if !response.status().is_success() {
        return Err(format!("HTTP request failed with status: {}", response.status()).into());
    }

    Ok(response)
}

// Process SSE response stream from the HTTP endpoint
async fn process_sse_response(response: reqwest::Response) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
    let mut stream = response.bytes_stream();
    let mut result_buffer = String::new();
    let mut final_result = String::new();
    let mut token_usage: Option<TokenUsage> = None;
    let mut confidence: Option<u8> = None;

    while let Some(chunk_result) = stream.next().await {
        let chunk = chunk_result?;
        let chunk_str = String::from_utf8_lossy(&chunk);

        for line in chunk_str.lines() {
            if let Some(data) = line.strip_prefix("data: ") {
                process_sse_event(
                    data,
                    &mut result_buffer,
                    &mut final_result,
                    &mut token_usage,
                    &mut confidence,
                )?;
            }
        }
    }

    Ok(build_complete_response(
        &result_buffer,
        &final_result,
        token_usage.as_ref(),
        confidence,
    ))
}

// Process individual SSE event
fn process_sse_event(
    data: &str,
    result_buffer: &mut String,
    final_result: &mut String,
    token_usage: &mut Option<TokenUsage>,
    confidence: &mut Option<u8>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    if let Ok(progress) = serde_json::from_str::<serde_json::Value>(data)
        && let Some(event_type) = progress.as_object().and_then(|obj| obj.keys().next())
    {
        match event_type.as_str() {
            "Status" => handle_status_event(&progress, result_buffer),
            "Schema" => handle_schema_event(result_buffer),
            "CypherQuery" => handle_cypher_query_event(&progress, result_buffer),
            "CypherResult" => handle_cypher_result_event(&progress, result_buffer),
            "ModelOutputChunk" => handle_model_output_chunk(&progress, final_result),
            "Result" => handle_result_event(&progress, final_result),
            "Confidence" => handle_confidence_event(&progress, confidence),
            "Usage" => handle_usage_event(&progress, token_usage),
            "Error" => return handle_error_event(&progress),
            _ => tracing::debug!("Unknown event type: {}", event_type),
        }
    }
    Ok(())
}

// Handle different types of SSE events
fn handle_status_event(
    progress: &serde_json::Value,
    result_buffer: &mut String,
) {
    if let Some(status) = progress.get("Status").and_then(|v| v.as_str()) {
        tracing::info!("Status: {}", status);
        writeln!(result_buffer, "Status: {status}").unwrap();
    }
}

fn handle_schema_event(result_buffer: &mut String) {
    tracing::info!("Schema discovered");
    result_buffer.push_str("Schema: Discovered\n");
}

fn handle_cypher_query_event(
    progress: &serde_json::Value,
    result_buffer: &mut String,
) {
    if let Some(query) = progress.get("CypherQuery").and_then(|v| v.as_str()) {
        tracing::info!("Generated Cypher: {}", query);
        writeln!(result_buffer, "Cypher Query: {query}").unwrap();
    }
}

fn handle_cypher_result_event(
    progress: &serde_json::Value,
    result_buffer: &mut String,
) {
    if let Some(cypher_result) = progress.get("CypherResult").and_then(|v| v.as_str()) {
        tracing::info!("Cypher result: {}", cypher_result);
        writeln!(result_buffer, "Query Result: {cypher_result}").unwrap();
    }
}

fn handle_model_output_chunk(
    progress: &serde_json::Value,
    final_result: &mut String,
) {
    if let Some(chunk) = progress.get("ModelOutputChunk").and_then(|v| v.as_str()) {
        final_result.push_str(chunk);
    }
}

fn handle_result_event(
    progress: &serde_json::Value,
    final_result: &mut String,
) {
    if let Some(result) = progress.get("Result").and_then(|v| v.as_str()) {
        tracing::info!("Final result received");
        *final_result = result.to_string();
    }
}

fn handle_confidence_event(
    progress: &serde_json::Value,
    confidence: &mut Option<u8>,
) {
    if let Some(value) = progress.get("Confidence").and_then(serde_json::Value::as_u64) {
        let value = u8::try_from(value.min(100)).unwrap_or(100);
        tracing::info!("Answer confidence: {}", value);
        *confidence = Some(value);
    }
}

fn handle_error_event(progress: &serde_json::Value) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    if let Some(error) = progress.get("Error").and_then(|v| v.as_str()) {
        tracing::error!("Error from HTTP endpoint: {}", error);
        return Err(format!("Error from text-to-cypher service: {error}").into());
    }
    Ok(())
}

fn handle_usage_event(
    progress: &serde_json::Value,
    token_usage: &mut Option<TokenUsage>,
) {
    if let Some(usage) = progress
        .get("Usage")
        .and_then(|v| serde_json::from_value::<TokenUsage>(v.clone()).ok())
    {
        tracing::info!(
            "Token usage: prompt={}, completion={}, total={}",
            usage.prompt_tokens,
            usage.completion_tokens,
            usage.total_tokens
        );
        *token_usage = Some(usage);
    }
}

// Build the complete response from buffer and final result
fn build_complete_response(
    result_buffer: &str,
    final_result: &str,
    token_usage: Option<&TokenUsage>,
    confidence: Option<u8>,
) -> String {
    let mut response = if final_result.is_empty() {
        result_buffer.trim().to_string()
    } else {
        format!("{}\n\nFinal Answer:\n{}", result_buffer.trim(), final_result)
    };

    if let Some(confidence) = confidence {
        write!(response, "\n\nConfidence: {confidence}%").unwrap();
    }

    if let Some(usage) = token_usage {
        write!(
            response,
            "\n\nToken Usage: prompt={}, completion={}, total={}",
            usage.prompt_tokens, usage.completion_tokens, usage.total_tokens
        )
        .unwrap();
    }

    response
}

// Helper function to get list of graphs from FalkorDB via REST API
async fn get_falkordb_graphs() -> Result<Vec<String>, Box<dyn std::error::Error + Send + Sync>> {
    // Call the local REST API endpoint
    let client = reqwest::Client::new();
    let response = client
        .get("http://localhost:8080/list_graphs")
        .send()
        .await
        .map_err(|e| format!("Failed to call list_graphs API: {e}"))?;

    if response.status().is_success() {
        let graphs: Vec<String> = response.json().await.map_err(|e| format!("Failed to parse response: {e}"))?;
        Ok(graphs)
    } else {
        Err(format!("API returned error status: {}", response.status()).into())
    }
}

// Helper function to get schema information for a specific graph via REST API
async fn get_graph_schema_via_api(graph_name: &str) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
    // Call the local REST API endpoint
    let client = reqwest::Client::new();
    let response = client
        .get(format!("http://localhost:8080/get_schema/{graph_name}"))
        .send()
        .await
        .map_err(|e| format!("Failed to call get_schema API: {e}"))?;

    if response.status().is_success() {
        let schema: String = response.json().await.map_err(|e| format!("Failed to parse response: {e}"))?;
        Ok(schema)
    } else {
        Err(format!("API returned error status: {}", response.status()).into())
    }
}

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

    // Feeds the SSE events a real run emits and returns the assembled MCP response.
    fn assemble(events: &[&str]) -> String {
        let mut result_buffer = String::new();
        let mut final_result = String::new();
        let mut token_usage: Option<TokenUsage> = None;
        let mut confidence: Option<u8> = None;
        for data in events {
            process_sse_event(
                data,
                &mut result_buffer,
                &mut final_result,
                &mut token_usage,
                &mut confidence,
            )
            .unwrap();
        }
        build_complete_response(&result_buffer, &final_result, token_usage.as_ref(), confidence)
    }

    #[test]
    fn confidence_event_is_surfaced_in_mcp_response() {
        let response = assemble(&[r#"{"Result":"The city names are A, B, C, and D."}"#, r#"{"Confidence":100}"#]);
        assert!(response.contains("Final Answer:\nThe city names are A, B, C, and D."));
        assert!(response.contains("Confidence: 100%"));
    }

    #[test]
    fn confidence_is_omitted_when_absent_and_clamped_when_high() {
        let without = assemble(&[r#"{"Result":"An answer."}"#]);
        assert!(!without.contains("Confidence:"));

        let clamped = assemble(&[r#"{"Result":"An answer."}"#, r#"{"Confidence":250}"#]);
        assert!(clamped.contains("Confidence: 100%"));
    }
}