service_kit 0.1.2

//! OpenAPI to MCP Router Builder

use crate::error::{Error, Result};
use crate::handler::ApiHandlerInventory;
use rmcp::handler::server::router::tool::{ToolRoute, ToolRouter};
use rmcp::model::{CallToolResult, Content, Tool};
use serde_json::{json, Map, Value};
use std::borrow::Cow;
use std::sync::Arc;
use utoipa::openapi::{OpenApi, PathItem, RefOr};
use utoipa::openapi::path::Operation;
use utoipa::openapi::schema::Schema;
use std::collections::HashMap;

pub type OutputSchemaTransformer = Arc<dyn Fn(&OpenApi, &Operation, &Option<Value>) -> Option<Value> + Send + Sync>;

#[derive(Default, Clone)]
pub struct OpenApiMcpRouterBuilder {
    openapi: Option<OpenApi>,
    output_transformer: Option<OutputSchemaTransformer>,
    router_state: Option<Arc<dyn std::any::Any + Send + Sync>>, 
}

impl OpenApiMcpRouterBuilder {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn openapi(mut self, openapi: OpenApi) -> Self {
        self.openapi = Some(openapi);
        self
    }

    pub fn with_output_transformer(mut self, transformer: OutputSchemaTransformer) -> Self {
        self.output_transformer = Some(transformer);
        self
    }

    /// 提供可选的 Router 状态（Any 透传给宏生成的执行函数）。
    pub fn with_state<S: Send + Sync + 'static>(mut self, state: S) -> Self {
        self.router_state = Some(Arc::new(state));
        self
    }

    pub fn build<S: Send + Sync + 'static>(self) -> Result<ToolRouter<S>> {
        let openapi = self.openapi.ok_or_else(|| {
            Error::SpecError("OpenAPI document not provided".to_string())
        })?;
        let handlers: HashMap<&'static str, for<'a> fn(&'a Value, axum::http::HeaderMap, Option<String>, Option<Arc<dyn std::any::Any + Send + Sync>>, Vec<u8>, Option<String>) -> crate::handler::DynHandlerFuture> =
            crate::inventory::iter::<ApiHandlerInventory>
                .into_iter()
                .map(|inv| (inv.operation_id, inv.handler))
                .collect();

        let mut router = ToolRouter::new();
        let router_state_outer = self.router_state.clone();

        for (_path, path_item) in openapi.paths.paths.iter() {
            for operation in operations_from_path_item(path_item) {
                if let Some(op_id) = operation.operation_id.as_deref() {
                    if let Some(handler_fn) = handlers.get(op_id).cloned() {
                        let tool_route = create_tool_route_for_handler(
                            &openapi,
                            (op_id.to_string(), handler_fn),
                            operation,
                            self.output_transformer.as_ref(),
                            router_state_outer.clone(),
                        )?;
                        router.add_route(tool_route);
                    }
                }
            }
        }

        Ok(router)
    }
}

fn create_tool_route_for_handler<S: Send + Sync + 'static>(
    openapi: &OpenApi,
    (operation_id, handler_fn): (
        String,
        for<'a> fn(&'a Value, axum::http::HeaderMap, Option<String>, Option<Arc<dyn std::any::Any + Send + Sync>>, Vec<u8>, Option<String>) -> crate::handler::DynHandlerFuture,
    ),
    operation: &Operation,
    output_transformer: Option<&OutputSchemaTransformer>,
    router_state: Option<Arc<dyn std::any::Any + Send + Sync>>,
) -> Result<ToolRoute<S>> {
    // Prefer requestBody schema; resolve $ref if present
    let body_schema_value: Option<Value> = operation
        .request_body
        .as_ref()
        .and_then(|body| body.content.get("application/json"))
        .and_then(|media_type| media_type.schema.as_ref())
        .and_then(|schema| resolve_schema(schema, openapi));

    // Aggregate parameters into a single object schema (query & path)
    let mut param_properties: Map<String, Value> = Map::new();
    let mut required_params: Vec<String> = Vec::new();
    if let Some(params) = operation.parameters.as_ref() {
        for p in params.iter() {
            let name: &str = &p.name;
            // 跳过认证相关字段
            if name.eq_ignore_ascii_case("authorization") { continue; }
            if let Some(schema) = p.schema.as_ref() {
                if let Some(schema_value) = resolve_schema(schema, openapi) {
                    param_properties.insert(name.to_string(), schema_value);
                }
            }
            if p.required == utoipa::openapi::Required::True {
                required_params.push(name.to_string());
            }
        }
    }

    let mut parameters_object_schema: Option<Value> = None;
    if !param_properties.is_empty() {
        let mut obj = Map::new();
        obj.insert("type".to_string(), Value::String("object".to_string()));
        obj.insert("properties".to_string(), Value::Object(param_properties.clone()));
        if !required_params.is_empty() {
            obj.insert(
                "required".to_string(),
                Value::Array(required_params.iter().cloned().map(Value::String).collect()),
            );
        }
        parameters_object_schema = Some(Value::Object(obj));
    }

    // Merge parameters into body schema when possible; prefer body on conflicts
    let mut input_schema = if let Some(body) = body_schema_value {
        let body_obj_opt: Option<Map<String, Value>> = body.as_object().cloned();
        let param_obj_opt: Option<Map<String, Value>> = parameters_object_schema
            .as_ref()
            .and_then(|v| v.as_object().cloned());
        if let (Some(mut body_obj), Some(param_obj)) = (body_obj_opt, param_obj_opt) {
            // Merge properties
            if let Some(Value::Object(mut body_props)) = body_obj.remove("properties") {
                if let Some(Value::Object(param_props)) = param_obj.get("properties") {
                    for (k, v) in param_props.iter() {
                        body_props.entry(k.clone()).or_insert(v.clone());
                    }
                }
                body_obj.insert("properties".to_string(), Value::Object(body_props));
            }
            // Merge required
            if let Some(Value::Array(mut body_required)) = body_obj.remove("required") {
                if let Some(Value::Array(param_required)) = param_obj.get("required") {
                    for v in param_required.iter() {
                        if !body_required.contains(v) {
                            body_required.push(v.clone());
                        }
                    }
                }
                body_obj.insert("required".to_string(), Value::Array(body_required));
            }
            Value::Object(body_obj)
        } else {
            body
        }
    } else if let Some(param_schema) = parameters_object_schema {
        param_schema
    } else {
        json!({ "type": "object" })
    };

    // Normalize schema to remove $ref/oneOf/additionalProperties for client compatibility
    input_schema = normalize_schema_value(&input_schema, openapi);

    let input_schema_map = if let Value::Object(map) = input_schema {
        Arc::new(map)
    } else {
        Arc::new(Map::new())
    };

    // Try resolving a 2xx application/json response schema as output_schema (prefer 200)
    let mut output_schema_value: Option<Value> = None;
    if let Some(resp) = operation.responses.responses.get("200").and_then(|r| match r {
        RefOr::T(r) => r.content.get("application/json"),
        _ => None,
    }) {
        if let Some(schema) = resp.schema.as_ref() { output_schema_value = resolve_schema(schema, openapi); }
    }
    if output_schema_value.is_none() {
        for (code, resp) in operation.responses.responses.iter() {
            if let Ok(code_num) = code.parse::<u16>() {
                if (200..300).contains(&code_num) {
                    if let RefOr::T(r) = resp {
                        if let Some(mt) = r.content.get("application/json") {
                            if let Some(schema) = mt.schema.as_ref() {
                                output_schema_value = resolve_schema(schema, openapi);
                                if output_schema_value.is_some() { break; }
                            }
                        }
                    }
                }
            }
        }
    }
    let mut output_schema_value = output_schema_value.map(|v| normalize_schema_value(&v, openapi));

    // 解包 ApiResponse<T>：若检测到外层是标准响应（含 success/timestamp 且有 data），提取 data 的内部 T
    if let Some(Value::Object(obj)) = &output_schema_value {
        if let Some(Value::Object(props)) = obj.get("properties") {
            let has_success = props.contains_key("success");
            let has_timestamp = props.contains_key("timestamp");
            if has_success && has_timestamp {
                if let Some(data_schema) = props.get("data") {
                    // 支持 oneOf [null, X]；$ref；或已具体化对象
                    let candidate: Option<Value> = match data_schema {
                        Value::Object(m) => {
                            if let Some(Value::Array(list)) = m.get("oneOf") {
                                list.iter().find(|v| {
                                    !(v.get("type").and_then(|t| t.as_str()).unwrap_or("") == "null")
                                }).cloned()
                            } else if m.contains_key("$ref") {
                                Some(Value::Object(m.clone()))
                            } else if m.contains_key("type") || m.contains_key("properties") || m.contains_key("items") {
                                Some(Value::Object(m.clone()))
                            } else {
                                None
                            }
                        }
                        other => Some(other.clone()),
                    };
                    if let Some(u) = candidate {
                        output_schema_value = Some(normalize_schema_value(&u, openapi));
                    }
                }
            }
        }
    }

    // 不在通用库内做与具体系统相关的 operationId → 组件映射回退，由项目侧 transformer 负责

    // 应用外部提供的 output schema 变换（例如项目侧定义的 ApiDefinition<T> 深度解包策略）
    if let Some(transformer) = output_transformer { output_schema_value = transformer(openapi, operation, &output_schema_value); }

    let output_schema_map = output_schema_value.and_then(|v| v.as_object().cloned()).map(Arc::new);
    let expects_structured_output = output_schema_map.is_some();

    let tool_def = Tool {
        name: operation_id.clone().into(),
        title: operation.summary.clone(),
        description: operation.description.clone().map(Cow::from),
        input_schema: input_schema_map,
        output_schema: output_schema_map,
        annotations: Default::default(),
        icons: None,
    };

    let route = ToolRoute::new_dyn(tool_def, move |ctx| {
        let handler_clone = handler_fn;
        let router_state_for_call = router_state.clone();
        Box::pin(async move {
            let params = ctx
                .arguments
                .as_ref()
                .map(|v| Value::Object(v.clone()))
                .unwrap_or(json!({}));

            // 从 MCP 上下文扩展中提取原始 HTTP 请求头，若不存在则为空
            let headers = ctx
                .request_context
                .extensions
                .get::<axum::http::request::Parts>()
                .map(|parts| parts.headers.clone())
                .unwrap_or_else(axum::http::HeaderMap::new);

            // 将 MCP 参数作为 JSON 负载传递给通用 REST 处理器：
            let body_bytes = serde_json::to_vec(&params).unwrap_or_default();
            let content_type = Some("application/json".to_string());
            // 标记来自 MCP 的桥接调用
            let mut headers = headers;
            headers.insert(axum::http::header::HeaderName::from_static("x-from-mcp"), axum::http::HeaderValue::from_static("1"));
            match handler_clone(&params, headers, None, router_state_for_call.clone(), body_bytes, content_type).await {
                Ok(response) => {
                    let (parts, body) = response.into_parts();
                    let body_bytes =
                        axum::body::to_bytes(body, usize::MAX).await.unwrap_or_default();

                    // 优先处理非 2xx 响应为错误
                    if !parts.status.is_success() {
                        let body_str = String::from_utf8_lossy(&body_bytes).to_string();
                        let err_msg =
                            format!("Handler failed with status {}: {}", parts.status, body_str);
                        return Ok(CallToolResult::error(vec![Content::text(err_msg)]));
                    }

                    // 即使 HTTP 成功，也检查业务层是否显示失败
                    let body_value: serde_json::Value = serde_json::from_slice(&body_bytes)
                        .unwrap_or(serde_json::Value::Null);
                    if let Some(success_flag) = body_value
                        .get("success")
                        .and_then(|v| v.as_bool())
                    {
                        if !success_flag {
                            let err_msg = format!(
                                "Handler reported business failure: {}",
                                body_value.to_string()
                            );
                            return Ok(CallToolResult::error(vec![Content::text(err_msg)]));
                        }
                    }

                    // 业务与 HTTP 均成功，按 output_schema 决定返回结构
                    let body_str = String::from_utf8_lossy(&body_bytes).to_string();
                    if expects_structured_output {
                        // 当工具声明了 output_schema，必须返回 structured_content
                        match serde_json::from_str::<serde_json::Value>(&body_str) {
                            Ok(json_val) => {
                                Ok(CallToolResult {
                                    content: Vec::new(),
                                    structured_content: Some(json_val),
                                    is_error: Some(false),
                                    meta: None,
                                })
                            }
                            // 如果服务端返回的不是 JSON，但 schema 存在，仍然尝试用文本回退为结构化包装
                            Err(_) => {
                                let wrapped = serde_json::json!({ "text": body_str });
                                Ok(CallToolResult {
                                    content: Vec::new(),
                                    structured_content: Some(wrapped),
                                    is_error: Some(false),
                                    meta: None,
                                })
                            }
                        }
                    } else {
                        Ok(CallToolResult::success(vec![Content::text(body_str)]))
                    }
                }
                Err(e) => {
                    let err_msg = format!("Handler execution failed: {}", e);
                    Ok(CallToolResult::error(vec![Content::text(err_msg)]))
                }
            }
        })
    });

    Ok(route)
}

fn operations_from_path_item(path_item: &PathItem) -> Vec<&Operation> {
    [
        &path_item.get,
        &path_item.post,
        &path_item.put,
        &path_item.delete,
        &path_item.options,
        &path_item.head,
        &path_item.patch,
        &path_item.trace,
    ]
    .iter()
    .filter_map(|op| op.as_ref())
    .collect()
}

/// Recursively normalize JSON schema value:
/// - Resolve top-level and nested $ref using OpenAPI components
/// - Replace oneOf [null, X] with X (optionality is captured by `required` at object level)
/// - Drop unsupported fields like additionalProperties/propertyNames for simpler clients
/// - Normalize nested properties and items
pub fn normalize_schema_value(value: &Value, openapi: &OpenApi) -> Value {
    match value {
        Value::Object(obj) => {
            // $ref handling
            if let Some(Value::String(reference)) = obj.get("$ref") {
                // Resolve and normalize referenced schema
                let name_opt = reference.rsplit('/').next();
                if let Some(name) = name_opt {
                    if let Some(components) = openapi.components.as_ref() {
                        if let Some(referenced) = components.schemas.get(name) {
                            match referenced {
                                RefOr::T(s) => {
                                    if let Ok(val) = serde_json::to_value(s) {
                                        return normalize_schema_value(&val, openapi);
                                    }
                                }
                                RefOr::Ref(r) => {
                                    let ref_val = json!({"$ref": r.ref_location.clone()});
                                    return normalize_schema_value(&ref_val, openapi);
                                }
                            }
                        }
                    }
                }
            }

            // type: ["null", "T"] handling → collapse to non-null type
            if let Some(Value::Array(type_list)) = obj.get("type") {
                let non_null = type_list.iter().find_map(|v| v.as_str().filter(|s| *s != "null"));
                if let Some(t) = non_null {
                    let mut out = obj.clone();
                    out.insert("type".into(), Value::String(t.to_string()));
                    return normalize_schema_value(&Value::Object(out), openapi);
                }
            }

            // oneOf handling: prefer first non-null schema
            if let Some(Value::Array(one_of_list)) = obj.get("oneOf") {
                for candidate in one_of_list {
                    let is_null = candidate
                        .get("type")
                        .and_then(|t| t.as_str())
                        .map(|t| t == "null")
                        .unwrap_or(false);
                    if !is_null {
                        return normalize_schema_value(candidate, openapi);
                    }
                }
            }

            // allOf/anyOf basic handling: take first element
            if let Some(Value::Array(all_of)) = obj.get("allOf") {
                if let Some(first) = all_of.first() {
                    return normalize_schema_value(first, openapi);
                }
            }
            if let Some(Value::Array(any_of)) = obj.get("anyOf") {
                if let Some(first) = any_of.first() {
                    return normalize_schema_value(first, openapi);
                }
            }

            // Clone and clean object
            let mut out = obj.clone();
            out.remove("$ref");
            out.remove("oneOf");
            out.remove("allOf");
            out.remove("anyOf");
            out.remove("additionalProperties");
            out.remove("propertyNames");

            // Normalize items
            if let Some(items) = out.get("items").cloned() {
                out.insert("items".into(), normalize_schema_value(&items, openapi));
            }

            // Normalize properties
            if let Some(Value::Object(props)) = out.get("properties") {
                let mut new_props = Map::new();
                for (k, v) in props.iter() {
                    new_props.insert(k.clone(), normalize_schema_value(v, openapi));
                }
                out.insert("properties".into(), Value::Object(new_props));
            }

            Value::Object(out)
        }
        Value::Array(arr) => {
            Value::Array(arr.iter().map(|v| normalize_schema_value(v, openapi)).collect())
        }
        _ => value.clone(),
    }
}

/// Resolve a RefOr<Schema> into a JSON value by following $ref using the OpenAPI components
fn resolve_schema(schema: &RefOr<Schema>, openapi: &OpenApi) -> Option<Value> {
    match schema {
        RefOr::T(s) => serde_json::to_value(s).ok(),
        RefOr::Ref(reference) => {
            // Expect format like "#/components/schemas/TypeName"
            let reference_str = reference.ref_location.clone();
            let name = reference_str.rsplit('/').next()?;
            let components = openapi.components.as_ref()?;
            let referenced = components.schemas.get(name)?;
            match referenced {
                RefOr::T(s) => serde_json::to_value(s).ok(),
                // Avoid deep recursion for now; return None if it's an indirect ref
                RefOr::Ref(_) => None,
            }
        }
    }
}