cc-switch/src-tauri/src/usage_script.rs

use reqwest::Client;
use rquickjs::{Context, Function, Runtime};
use serde_json::Value;
use std::collections::HashMap;
use std::time::Duration;

use crate::error::AppError;

/// 执行用量查询脚本
pub async fn execute_usage_script(
    script_code: &str,
    api_key: &str,
    base_url: &str,
    timeout_secs: u64,
    access_token: Option<&str>,
    user_id: Option<&str>,
) -> Result<Value, AppError> {
    // 1. 替换变量
    let mut replaced = script_code
        .replace("{{apiKey}}", api_key)
        .replace("{{baseUrl}}", base_url);

    // 替换 accessToken 和 userId
    if let Some(token) = access_token {
        replaced = replaced.replace("{{accessToken}}", token);
    }
    if let Some(uid) = user_id {
        replaced = replaced.replace("{{userId}}", uid);
    }

    // 2. 在独立作用域中提取 request 配置（确保 Runtime/Context 在 await 前释放）
    let request_config = {
        let runtime =
            Runtime::new().map_err(|e| AppError::localized("usage_script.runtime_create_failed", format!("创建 JS 运行时失败: {}", e), format!("Failed to create JS runtime: {}", e)))?;
        let context = Context::full(&runtime)
            .map_err(|e| AppError::localized("usage_script.context_create_failed", format!("创建 JS 上下文失败: {}", e), format!("Failed to create JS context: {}", e)))?;

        context.with(|ctx| {
            // 执行用户代码，获取配置对象
            let config: rquickjs::Object = ctx
                .eval(replaced.clone())
                .map_err(|e| AppError::localized("usage_script.config_parse_failed", format!("解析配置失败: {}", e), format!("Failed to parse config: {}", e)))?;

            // 提取 request 配置
            let request: rquickjs::Object = config
                .get("request")
                .map_err(|e| AppError::localized("usage_script.request_missing", format!("缺少 request 配置: {}", e), format!("Missing request config: {}", e)))?;

            // 将 request 转换为 JSON 字符串
            let request_json: String = ctx
                .json_stringify(request)
                .map_err(|e| AppError::localized("usage_script.request_serialize_failed", format!("序列化 request 失败: {}", e), format!("Failed to serialize request: {}", e)))?
                .ok_or_else(|| AppError::localized("usage_script.serialize_none", "序列化返回 None", "Serialization returned None"))?
                .get()
                .map_err(|e| AppError::localized("usage_script.get_string_failed", format!("获取字符串失败: {}", e), format!("Failed to get string: {}", e)))?;

            Ok::<_, AppError>(request_json)
        })?
    }; // Runtime 和 Context 在这里被 drop

    // 3. 解析 request 配置
    let request: RequestConfig = serde_json::from_str(&request_config)
        .map_err(|e| AppError::localized("usage_script.request_format_invalid", format!("request 配置格式错误: {}", e), format!("Invalid request config format: {}", e)))?;

    // 4. 发送 HTTP 请求
    let response_data = send_http_request(&request, timeout_secs).await?;

    // 5. 在独立作用域中执行 extractor（确保 Runtime/Context 在函数结束前释放）
    let result: Value = {
        let runtime =
            Runtime::new().map_err(|e| AppError::localized("usage_script.runtime_create_failed", format!("创建 JS 运行时失败: {}", e), format!("Failed to create JS runtime: {}", e)))?;
        let context = Context::full(&runtime)
            .map_err(|e| AppError::localized("usage_script.context_create_failed", format!("创建 JS 上下文失败: {}", e), format!("Failed to create JS context: {}", e)))?;

        context.with(|ctx| {
            // 重新 eval 获取配置对象
            let config: rquickjs::Object = ctx
                .eval(replaced.clone())
                .map_err(|e| AppError::localized("usage_script.config_reparse_failed", format!("重新解析配置失败: {}", e), format!("Failed to re-parse config: {}", e)))?;

            // 提取 extractor 函数
            let extractor: Function = config
                .get("extractor")
                .map_err(|e| AppError::localized("usage_script.extractor_missing", format!("缺少 extractor 函数: {}", e), format!("Missing extractor function: {}", e)))?;

            // 将响应数据转换为 JS 值
            let response_js: rquickjs::Value = ctx
                .json_parse(response_data.as_str())
                .map_err(|e| AppError::localized("usage_script.response_parse_failed", format!("解析响应 JSON 失败: {}", e), format!("Failed to parse response JSON: {}", e)))?;

            // 调用 extractor(response)
            let result_js: rquickjs::Value = extractor
                .call((response_js,))
                .map_err(|e| AppError::localized("usage_script.extractor_exec_failed", format!("执行 extractor 失败: {}", e), format!("Failed to execute extractor: {}", e)))?;

            // 转换为 JSON 字符串
            let result_json: String = ctx
                .json_stringify(result_js)
                .map_err(|e| AppError::localized("usage_script.result_serialize_failed", format!("序列化结果失败: {}", e), format!("Failed to serialize result: {}", e)))?
                .ok_or_else(|| AppError::localized("usage_script.serialize_none", "序列化返回 None", "Serialization returned None"))?
                .get()
                .map_err(|e| AppError::localized("usage_script.get_string_failed", format!("获取字符串失败: {}", e), format!("Failed to get string: {}", e)))?;

            // 解析为 serde_json::Value
            serde_json::from_str(&result_json)
                .map_err(|e| AppError::localized("usage_script.json_parse_failed", format!("JSON 解析失败: {}", e), format!("JSON parse failed: {}", e)))
        })?
    }; // Runtime 和 Context 在这里被 drop

    // 6. 验证返回值格式
    validate_result(&result)?;

    Ok(result)
}

/// 请求配置结构
#[derive(Debug, serde::Deserialize)]
struct RequestConfig {
    url: String,
    method: String,
    #[serde(default)]
    headers: HashMap<String, String>,
    #[serde(default)]
    body: Option<String>,
}

/// 发送 HTTP 请求
async fn send_http_request(config: &RequestConfig, timeout_secs: u64) -> Result<String, AppError> {
    // 约束超时范围，防止异常配置导致长时间阻塞
    let timeout = timeout_secs.clamp(2, 30);
    let client = Client::builder()
        .timeout(Duration::from_secs(timeout))
        .build()
        .map_err(|e| AppError::localized("usage_script.client_create_failed", format!("创建客户端失败: {}", e), format!("Failed to create client: {}", e)))?;

    // 严格校验 HTTP 方法，非法值不回退为 GET
    let method: reqwest::Method = config
        .method
        .parse()
        .map_err(|_| AppError::InvalidHttpMethod(config.method.clone()))?;

    let mut req = client.request(method.clone(), &config.url);

    // 添加请求头
    for (k, v) in &config.headers {
        req = req.header(k, v);
    }

    // 添加请求体
    if let Some(body) = &config.body {
        req = req.body(body.clone());
    }

    // 发送请求
    let resp = req
        .send()
        .await
        .map_err(|e| AppError::localized("usage_script.request_failed", format!("请求失败: {}", e), format!("Request failed: {}", e)))?;

    let status = resp.status();
    let text = resp
        .text()
        .await
        .map_err(|e| AppError::localized("usage_script.read_response_failed", format!("读取响应失败: {}", e), format!("Failed to read response: {}", e)))?;

    if !status.is_success() {
        let preview = if text.len() > 200 {
            format!("{}...", &text[..200])
        } else {
            text.clone()
        };
        return Err(AppError::localized("usage_script.http_error", format!("HTTP {} : {}", status, preview), format!("HTTP {} : {}", status, preview)));
    }

    Ok(text)
}

/// 验证脚本返回值（支持单对象或数组）
fn validate_result(result: &Value) -> Result<(), AppError> {
    // 如果是数组，验证每个元素
    if let Some(arr) = result.as_array() {
        if arr.is_empty() {
            return Err(AppError::localized("usage_script.empty_array", "脚本返回的数组不能为空", "Script returned empty array"));
        }
        for (idx, item) in arr.iter().enumerate() {
            validate_single_usage(item)
                .map_err(|e| AppError::localized("usage_script.array_validation_failed", format!("数组索引[{}]验证失败: {}", idx, e), format!("Validation failed at index [{}]: {}", idx, e)))?;
        }
        return Ok(());
    }

    // 如果是单对象，直接验证（向后兼容）
    validate_single_usage(result)
}

/// 验证单个用量数据对象
fn validate_single_usage(result: &Value) -> Result<(), AppError> {
    let obj = result
        .as_object()
        .ok_or_else(|| AppError::localized("usage_script.must_return_object", "脚本必须返回对象或对象数组", "Script must return object or array of objects"))?;

    // 所有字段均为可选，只进行类型检查
    if obj.contains_key("isValid")
        && !result["isValid"].is_null()
        && !result["isValid"].is_boolean()
    {
        return Err(AppError::localized("usage_script.isvalid_type_error", "isValid 必须是布尔值或 null", "isValid must be boolean or null"));
    }
    if obj.contains_key("invalidMessage")
        && !result["invalidMessage"].is_null()
        && !result["invalidMessage"].is_string()
    {
        return Err(AppError::localized("usage_script.invalidmessage_type_error", "invalidMessage 必须是字符串或 null", "invalidMessage must be string or null"));
    }
    if obj.contains_key("remaining")
        && !result["remaining"].is_null()
        && !result["remaining"].is_number()
    {
        return Err(AppError::localized("usage_script.remaining_type_error", "remaining 必须是数字或 null", "remaining must be number or null"));
    }
    if obj.contains_key("unit") && !result["unit"].is_null() && !result["unit"].is_string() {
        return Err(AppError::localized("usage_script.unit_type_error", "unit 必须是字符串或 null", "unit must be string or null"));
    }
    if obj.contains_key("total") && !result["total"].is_null() && !result["total"].is_number() {
        return Err(AppError::localized("usage_script.total_type_error", "total 必须是数字或 null", "total must be number or null"));
    }
    if obj.contains_key("used") && !result["used"].is_null() && !result["used"].is_number() {
        return Err(AppError::localized("usage_script.used_type_error", "used 必须是数字或 null", "used must be number or null"));
    }
    if obj.contains_key("planName")
        && !result["planName"].is_null()
        && !result["planName"].is_string()
    {
        return Err(AppError::localized("usage_script.planname_type_error", "planName 必须是字符串或 null", "planName must be string or null"));
    }
    if obj.contains_key("extra") && !result["extra"].is_null() && !result["extra"].is_string() {
        return Err(AppError::localized("usage_script.extra_type_error", "extra 必须是字符串或 null", "extra must be string or null"));
    }

    Ok(())
}