mod app;
mod catalog;
mod jsonrpc;
mod session;

use std::{env, net::SocketAddr, time::Duration};

use mcpaas_registry::PostgresRegistry;
use tokio::net::TcpListener;
use tracing::info;

use crate::app::build_app;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    tracing_subscriber::fmt()
        .with_env_filter(
            env::var("MCPAAS_LOG_LEVEL")
                .unwrap_or_else(|_| "mcp_server=info,tower_http=info".into()),
        )
        .init();

    let database_url = env::var("MCPAAS_DATABASE_URL")?;
    let bind_addr = env::var("MCPAAS_MCP_BIND").unwrap_or_else(|_| "0.0.0.0:3002".into());
    let public_base_url = env::var("MCPAAS_PUBLIC_BASE_URL").ok();
    let refresh_interval = env::var("MCPAAS_MCP_REFRESH_MS")
        .ok()
        .and_then(|value| value.parse::<u64>().ok())
        .map(Duration::from_millis)
        .unwrap_or_else(|| Duration::from_secs(5));
    let socket_addr: SocketAddr = bind_addr.parse()?;
    let registry = PostgresRegistry::connect(&database_url).await?;
    let app = build_app(registry, refresh_interval, public_base_url);
    let listener = TcpListener::bind(socket_addr).await?;

    info!("mcp-server listening on {}", socket_addr);

    axum::serve(listener, app).await?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use std::{
        collections::BTreeMap,
        env,
        time::{Duration, SystemTime, UNIX_EPOCH},
    };

    use axum::{Json, Router, http::header, routing::post};
    use mcpaas_adapter_grpc::test_support as grpc_test_support;
    use mcpaas_core::{
        DescriptorId, ExecutionConfig, GraphqlOperationType, GraphqlTarget, GrpcTarget, HttpMethod,
        Operation, OperationId, OperationStatus, Protocol, RestTarget, Target, ToolDescription,
    };
    use mcpaas_mapping::{MappingRule, MappingSet};
    use mcpaas_registry::{PostgresRegistry, PublishRequest};
    use mcpaas_schema::{Schema, SchemaKind};
    use serde_json::{Value, json};
    use sqlx::{Executor, postgres::PgPoolOptions};
    use tokio::net::TcpListener;

    use crate::app::build_app;

    #[tokio::test]
    async fn initializes_lists_and_calls_published_tool_via_mcp() {
        let registry = test_registry().await;
        let upstream_base_url = spawn_upstream_server().await;
        let operation = test_operation(&upstream_base_url, "crm_create_lead");

        registry
            .create_operation(&operation, Some("alice"))
            .await
            .unwrap();
        registry
            .publish_operation(PublishRequest {
                operation_id: &operation.id,
                version: 1,
                published_at: "2026-03-26T10:00:00Z",
                published_by: Some("alice"),
            })
            .await
            .unwrap();

        let base_url = spawn_mcp_server(build_app(
            registry,
            Duration::from_millis(0),
            Some("https://rmcp.example.com".to_owned()),
        ))
        .await;
        let client = reqwest::Client::new();
        let initialized_session = initialize_session(&client, &base_url).await;

        let tools = post_jsonrpc(
            &client,
            &base_url,
            Some(&initialized_session),
            json!({
                "jsonrpc": "2.0",
                "id": 2,
                "method": "tools/list",
                "params": {}
            }),
        )
        .await;
        let call_result = post_jsonrpc(
            &client,
            &base_url,
            Some(&initialized_session),
            json!({
                "jsonrpc": "2.0",
                "id": 3,
                "method": "tools/call",
                "params": {
                    "name": "crm_create_lead",
                    "arguments": {
                        "email": "user@example.com"
                    }
                }
            }),
        )
        .await;

        assert_eq!(tools["result"]["tools"][0]["name"], "crm_create_lead");
        assert_eq!(
            call_result["result"]["structuredContent"],
            json!({ "id": "lead_123" })
        );
        assert_eq!(call_result["result"]["isError"], false);
    }

    #[tokio::test]
    async fn initializes_and_calls_published_graphql_tool_via_mcp() {
        let registry = test_registry().await;
        let endpoint = spawn_graphql_server().await;
        let operation = test_graphql_operation(&endpoint, "crm_create_lead_graphql");

        registry
            .create_operation(&operation, Some("alice"))
            .await
            .unwrap();
        registry
            .publish_operation(PublishRequest {
                operation_id: &operation.id,
                version: 1,
                published_at: "2026-03-26T10:00:00Z",
                published_by: Some("alice"),
            })
            .await
            .unwrap();

        let base_url = spawn_mcp_server(build_app(
            registry,
            Duration::from_millis(0),
            Some("https://rmcp.example.com".to_owned()),
        ))
        .await;
        let client = reqwest::Client::new();
        let initialized_session = initialize_session(&client, &base_url).await;

        let call_result = post_jsonrpc(
            &client,
            &base_url,
            Some(&initialized_session),
            json!({
                "jsonrpc": "2.0",
                "id": 3,
                "method": "tools/call",
                "params": {
                    "name": "crm_create_lead_graphql",
                    "arguments": {
                        "email": "user@example.com"
                    }
                }
            }),
        )
        .await;

        assert_eq!(
            call_result["result"]["structuredContent"],
            json!({ "id": "lead_123" })
        );
        assert_eq!(call_result["result"]["isError"], false);
    }

    #[tokio::test]
    async fn initializes_and_calls_published_grpc_tool_via_mcp() {
        let registry = test_registry().await;
        let server_addr = grpc_test_support::spawn_unary_echo_server().await;
        let operation = test_grpc_operation(&server_addr, "echo_unary_grpc");

        registry
            .create_operation(&operation, Some("alice"))
            .await
            .unwrap();
        registry
            .publish_operation(PublishRequest {
                operation_id: &operation.id,
                version: 1,
                published_at: "2026-03-26T10:00:00Z",
                published_by: Some("alice"),
            })
            .await
            .unwrap();

        let base_url = spawn_mcp_server(build_app(
            registry,
            Duration::from_millis(0),
            Some("https://rmcp.example.com".to_owned()),
        ))
        .await;
        let client = reqwest::Client::new();
        let initialized_session = initialize_session(&client, &base_url).await;

        let call_result = post_jsonrpc(
            &client,
            &base_url,
            Some(&initialized_session),
            json!({
                "jsonrpc": "2.0",
                "id": 3,
                "method": "tools/call",
                "params": {
                    "name": "echo_unary_grpc",
                    "arguments": {
                        "message": "hello"
                    }
                }
            }),
        )
        .await;

        assert_eq!(
            call_result["result"]["structuredContent"],
            json!({ "message": "hello" })
        );
        assert_eq!(call_result["result"]["isError"], false);
    }

    #[tokio::test]
    async fn requires_initialized_notification_before_tool_methods() {
        let registry = test_registry().await;
        let base_url = spawn_mcp_server(build_app(
            registry,
            Duration::from_millis(0),
            Some("https://rmcp.example.com".to_owned()),
        ))
        .await;
        let client = reqwest::Client::new();
        let initialize_response = client
            .post(format!("{base_url}/mcp"))
            .header(header::ACCEPT, "application/json, text/event-stream")
            .json(&json!({
                "jsonrpc": "2.0",
                "id": 1,
                "method": "initialize",
                "params": {
                    "protocolVersion": "2025-11-25"
                }
            }))
            .send()
            .await
            .unwrap();
        let session_id = initialize_response
            .headers()
            .get("MCP-Session-Id")
            .unwrap()
            .to_str()
            .unwrap()
            .to_owned();
        let tools_list = client
            .post(format!("{base_url}/mcp"))
            .header(header::ACCEPT, "application/json, text/event-stream")
            .header("MCP-Session-Id", session_id)
            .header("MCP-Protocol-Version", "2025-11-25")
            .json(&json!({
                "jsonrpc": "2.0",
                "id": 2,
                "method": "tools/list",
                "params": {}
            }))
            .send()
            .await
            .unwrap()
            .json::<Value>()
            .await
            .unwrap();

        assert_eq!(tools_list["error"]["code"], -32002);
    }

    #[tokio::test]
    async fn refreshes_published_tools_without_restart() {
        let registry = test_registry().await;
        let upstream_base_url = spawn_upstream_server().await;
        let base_url = spawn_mcp_server(build_app(
            registry.clone(),
            Duration::from_millis(0),
            Some("https://rmcp.example.com".to_owned()),
        ))
        .await;
        let client = reqwest::Client::new();
        let initialized_session = initialize_session(&client, &base_url).await;

        let before_publish = post_jsonrpc(
            &client,
            &base_url,
            Some(&initialized_session),
            json!({
                "jsonrpc": "2.0",
                "id": 2,
                "method": "tools/list",
                "params": {}
            }),
        )
        .await;

        let operation = test_operation(&upstream_base_url, "crm_publish_later");
        registry
            .create_operation(&operation, Some("alice"))
            .await
            .unwrap();
        registry
            .publish_operation(PublishRequest {
                operation_id: &operation.id,
                version: 1,
                published_at: "2026-03-26T10:00:00Z",
                published_by: Some("alice"),
            })
            .await
            .unwrap();

        let after_publish = post_jsonrpc(
            &client,
            &base_url,
            Some(&initialized_session),
            json!({
                "jsonrpc": "2.0",
                "id": 3,
                "method": "tools/list",
                "params": {}
            }),
        )
        .await;

        assert_eq!(before_publish["result"]["tools"], json!([]));
        assert_eq!(
            after_publish["result"]["tools"][0]["name"],
            "crm_publish_later"
        );
    }

    async fn initialize_session(client: &reqwest::Client, base_url: &str) -> String {
        let initialize_response = client
            .post(format!("{base_url}/mcp"))
            .header(header::ACCEPT, "application/json, text/event-stream")
            .json(&json!({
                "jsonrpc": "2.0",
                "id": 1,
                "method": "initialize",
                "params": {
                    "protocolVersion": "2025-11-25"
                }
            }))
            .send()
            .await
            .unwrap();
        let session_id = initialize_response
            .headers()
            .get("MCP-Session-Id")
            .unwrap()
            .to_str()
            .unwrap()
            .to_owned();

        let initialized_response = client
            .post(format!("{base_url}/mcp"))
            .header(header::ACCEPT, "application/json, text/event-stream")
            .header("MCP-Session-Id", &session_id)
            .header("MCP-Protocol-Version", "2025-11-25")
            .json(&json!({
                "jsonrpc": "2.0",
                "method": "notifications/initialized",
                "params": {}
            }))
            .send()
            .await
            .unwrap();

        assert_eq!(initialized_response.status(), reqwest::StatusCode::ACCEPTED);

        session_id
    }

    async fn post_jsonrpc(
        client: &reqwest::Client,
        base_url: &str,
        session_id: Option<&str>,
        payload: Value,
    ) -> Value {
        let mut request = client
            .post(format!("{base_url}/mcp"))
            .header(header::ACCEPT, "application/json, text/event-stream")
            .header("MCP-Protocol-Version", "2025-11-25");

        if let Some(session_id) = session_id {
            request = request.header("MCP-Session-Id", session_id);
        }

        request
            .json(&payload)
            .send()
            .await
            .unwrap()
            .json::<Value>()
            .await
            .unwrap()
    }

    async fn spawn_upstream_server() -> String {
        let app = Router::new().route("/crm/leads", post(create_lead));
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let address = listener.local_addr().unwrap();

        tokio::spawn(async move {
            axum::serve(listener, app).await.unwrap();
        });

        format!("http://{}", address)
    }

    async fn spawn_graphql_server() -> String {
        let app = Router::new().route("/", post(graphql_handler));
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let address = listener.local_addr().unwrap();

        tokio::spawn(async move {
            axum::serve(listener, app).await.unwrap();
        });

        format!("http://{}", address)
    }

    async fn spawn_mcp_server(app: Router) -> String {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let address = listener.local_addr().unwrap();

        tokio::spawn(async move {
            axum::serve(listener, app).await.unwrap();
        });

        format!("http://{}", address)
    }

    async fn create_lead(Json(payload): Json<Value>) -> Json<Value> {
        Json(json!({
            "id": "lead_123",
            "email": payload["email"]
        }))
    }

    async fn graphql_handler(Json(payload): Json<Value>) -> Json<Value> {
        let email = payload
            .get("variables")
            .and_then(|variables| variables.get("email"))
            .and_then(Value::as_str)
            .unwrap_or_default();

        Json(json!({
            "data": {
                "createLead": {
                    "id": "lead_123",
                    "email": email
                }
            }
        }))
    }

    async fn test_registry() -> PostgresRegistry {
        let database_url = env::var("TEST_DATABASE_URL")
            .unwrap_or_else(|_| "postgres://rmcp:rmcp@127.0.0.1:5432/rmcp".to_owned());
        let admin_pool = PgPoolOptions::new()
            .max_connections(1)
            .connect(&database_url)
            .await
            .unwrap();
        let schema = format!(
            "test_mcp_server_{}_{}",
            std::process::id(),
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        );

        admin_pool
            .execute(sqlx::query(&format!("create schema {schema}")))
            .await
            .unwrap();

        PostgresRegistry::connect(&format!("{database_url}?options=-csearch_path%3D{schema}"))
            .await
            .unwrap()
    }

    fn test_operation(base_url: &str, name: &str) -> Operation<Schema, MappingSet> {
        Operation {
            id: OperationId::new(format!("op_{name}")),
            name: name.to_owned(),
            display_name: "Create Lead".to_owned(),
            protocol: Protocol::Rest,
            status: OperationStatus::Published,
            version: 1,
            target: Target::Rest(RestTarget {
                base_url: base_url.to_owned(),
                method: HttpMethod::Post,
                path_template: "/crm/leads".to_owned(),
                static_headers: BTreeMap::new(),
            }),
            input_schema: object_schema("email"),
            output_schema: object_schema("id"),
            input_mapping: MappingSet {
                rules: vec![MappingRule {
                    source: "$.mcp.email".to_owned(),
                    target: "$.request.body.email".to_owned(),
                    required: true,
                    default_value: None,
                    transform: None,
                    condition: None,
                    notes: None,
                }],
            },
            output_mapping: MappingSet {
                rules: vec![MappingRule {
                    source: "$.response.body.id".to_owned(),
                    target: "$.output.id".to_owned(),
                    required: true,
                    default_value: None,
                    transform: None,
                    condition: None,
                    notes: None,
                }],
            },
            execution_config: ExecutionConfig {
                timeout_ms: 1_000,
                retry_policy: None,
                auth_profile_ref: None,
                headers: BTreeMap::new(),
                protocol_options: None,
            },
            tool_description: ToolDescription {
                title: "Create Lead".to_owned(),
                description: "Creates a CRM lead".to_owned(),
                tags: Vec::new(),
                examples: Vec::new(),
            },
            samples: None,
            generated_draft: None,
            config_export: None,
            created_at: "2026-03-26T10:00:00Z".to_owned(),
            updated_at: "2026-03-26T10:00:00Z".to_owned(),
            published_at: Some("2026-03-26T10:00:00Z".to_owned()),
        }
    }

    fn test_graphql_operation(endpoint: &str, name: &str) -> Operation<Schema, MappingSet> {
        Operation {
            id: OperationId::new(format!("op_{name}")),
            name: name.to_owned(),
            display_name: "Create Lead GraphQL".to_owned(),
            protocol: Protocol::Graphql,
            status: OperationStatus::Published,
            version: 1,
            target: Target::Graphql(GraphqlTarget {
                endpoint: endpoint.to_owned(),
                operation_type: GraphqlOperationType::Mutation,
                operation_name: "CreateLead".to_owned(),
                query_template:
                    "mutation CreateLead($email: String!) { createLead(email: $email) { id email } }"
                        .to_owned(),
                response_path: "$.response.body.data.createLead".to_owned(),
            }),
            input_schema: object_schema("email"),
            output_schema: object_schema("id"),
            input_mapping: MappingSet {
                rules: vec![MappingRule {
                    source: "$.mcp.email".to_owned(),
                    target: "$.request.variables.email".to_owned(),
                    required: true,
                    default_value: None,
                    transform: None,
                    condition: None,
                    notes: None,
                }],
            },
            output_mapping: MappingSet {
                rules: vec![MappingRule {
                    source: "$.response.data.id".to_owned(),
                    target: "$.output.id".to_owned(),
                    required: true,
                    default_value: None,
                    transform: None,
                    condition: None,
                    notes: None,
                }],
            },
            execution_config: ExecutionConfig {
                timeout_ms: 1_000,
                retry_policy: None,
                auth_profile_ref: None,
                headers: BTreeMap::new(),
                protocol_options: None,
            },
            tool_description: ToolDescription {
                title: "Create Lead GraphQL".to_owned(),
                description: "Creates a CRM lead through GraphQL".to_owned(),
                tags: vec!["graphql".to_owned()],
                examples: Vec::new(),
            },
            samples: None,
            generated_draft: None,
            config_export: None,
            created_at: "2026-03-26T10:00:00Z".to_owned(),
            updated_at: "2026-03-26T10:00:00Z".to_owned(),
            published_at: Some("2026-03-26T10:00:00Z".to_owned()),
        }
    }

    fn test_grpc_operation(server_addr: &str, name: &str) -> Operation<Schema, MappingSet> {
        Operation {
            id: OperationId::new(format!("op_{name}")),
            name: name.to_owned(),
            display_name: "Unary Echo gRPC".to_owned(),
            protocol: Protocol::Grpc,
            status: OperationStatus::Published,
            version: 1,
            target: Target::Grpc(GrpcTarget {
                server_addr: server_addr.to_owned(),
                package: "echo".to_owned(),
                service: "EchoService".to_owned(),
                method: "UnaryEcho".to_owned(),
                descriptor_ref: DescriptorId::new("desc_echo"),
                descriptor_set_b64: grpc_test_support::descriptor_set_b64(),
            }),
            input_schema: object_schema("message"),
            output_schema: object_schema("message"),
            input_mapping: MappingSet {
                rules: vec![MappingRule {
                    source: "$.mcp.message".to_owned(),
                    target: "$.request.grpc.message".to_owned(),
                    required: true,
                    default_value: None,
                    transform: None,
                    condition: None,
                    notes: None,
                }],
            },
            output_mapping: MappingSet {
                rules: vec![MappingRule {
                    source: "$.response.data.message".to_owned(),
                    target: "$.output.message".to_owned(),
                    required: true,
                    default_value: None,
                    transform: None,
                    condition: None,
                    notes: None,
                }],
            },
            execution_config: ExecutionConfig {
                timeout_ms: 1_000,
                retry_policy: None,
                auth_profile_ref: None,
                headers: BTreeMap::new(),
                protocol_options: None,
            },
            tool_description: ToolDescription {
                title: "Unary Echo gRPC".to_owned(),
                description: "Echoes a unary gRPC payload".to_owned(),
                tags: vec!["grpc".to_owned()],
                examples: Vec::new(),
            },
            samples: None,
            generated_draft: None,
            config_export: None,
            created_at: "2026-03-26T10:00:00Z".to_owned(),
            updated_at: "2026-03-26T10:00:00Z".to_owned(),
            published_at: Some("2026-03-26T10:00:00Z".to_owned()),
        }
    }

    fn object_schema(field_name: &str) -> Schema {
        Schema {
            kind: SchemaKind::Object,
            description: None,
            required: true,
            nullable: false,
            default_value: None,
            fields: BTreeMap::from([(
                field_name.to_owned(),
                Schema {
                    kind: SchemaKind::String,
                    description: None,
                    required: true,
                    nullable: false,
                    default_value: None,
                    fields: BTreeMap::new(),
                    items: None,
                    enum_values: Vec::new(),
                    variants: Vec::new(),
                },
            )]),
            items: None,
            enum_values: Vec::new(),
            variants: Vec::new(),
        }
    }
}