AstroResearch/src/agent/tools/astro/research/vizier.rs

// src/agent/tools/astro/research/vizier.rs
//
// CatalogOperationTool —— VizieR 星表统一操作工具
//   search: 按关键词搜索星表目录
//   describe: 查看表的列结构
//   query: 执行 ADQL 或按表名查询，返回数据
//   cone: 按坐标锥形检索
//   export: 导出为 CSV 文件
//   lookup: 通过 bibcode 查找关联数据表

use async_trait::async_trait;
use serde_json::json;
use tracing::info;

use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
use crate::clients::cds::vizier::VizierQueryResult;

/// 渲染查询结果为 Markdown 表格（前 N 行）
fn render_result_table(result: &VizierQueryResult, preview_rows: usize) -> String {
    let mut content = String::new();

    if let Some(ref table) = result.table_name {
        content.push_str(&format!("来源表: {}\n", table));
    }
    content.push_str(&format!(
        "共 {} 行（{}）\n\n",
        result.row_count,
        if result.truncated {
            "已截断，存在更多结果"
        } else {
            "完整结果"
        }
    ));

    if result.fields.is_empty() || result.rows.is_empty() {
        content.push_str("(无数据行)\n");
        return content;
    }

    let headers: Vec<&str> = result.fields.iter().map(|f| f.name.as_str()).collect();
    content.push_str(&format!("| {} |\n", headers.join(" | ")));
    content.push_str(&format!(
        "| {} |\n",
        headers
            .iter()
            .map(|_| "---")
            .collect::<Vec<_>>()
            .join(" | ")
    ));

    let show = result.rows.len().min(preview_rows);
    for row in result.rows.iter().take(show) {
        let cells: Vec<String> = row
            .iter()
            .map(|v| match v {
                serde_json::Value::Null => "—".to_string(),
                serde_json::Value::String(s) => s.clone(),
                other => other.to_string(),
            })
            .collect();
        content.push_str(&format!("| {} |\n", cells.join(" | ")));
    }

    if result.rows.len() > preview_rows {
        content.push_str(&format!(
            "\n(已省略 {} 行，完整数据见结构化输出)\n",
            result.rows.len() - preview_rows
        ));
    }

    let units: Vec<&str> = result
        .fields
        .iter()
        .filter_map(|f| f.unit.as_deref())
        .collect();
    if !units.is_empty() {
        content.push_str(&format!("\n字段单位: {}\n", units.join(", ")));
    }

    content
}

// ── CatalogOperationTool ──

pub struct CatalogOperationTool;

#[async_trait]
impl AgentTool for CatalogOperationTool {
    fn name(&self) -> &str {
        "catalog_operation"
    }

    fn display_name(&self) -> &str {
        "星表检索"
    }

    fn description(&self) -> &str {
        "VizieR 星表统一操作。通过 action 选择：\n\
         search — 按关键词搜索星表目录；\n\
         describe — 查看表的列结构；\n\
         query — 执行 ADQL 或按表名查询数据；\n\
         cone — 按坐标锥形检索近邻天体；\n\
         export — 下载表数据保存为 CSV；\n\
         lookup — 通过 bibcode 查找关联数据表。"
    }

    fn parameters(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "action": {
                    "type": "string",
                    "enum": ["search", "describe", "query", "cone", "export", "lookup"],
                    "description": "操作类型"
                },
                "keyword": {
                    "type": "string",
                    "description": "search 时的搜索关键词，如 'Gaia DR3'"
                },
                "table": {
                    "type": "string",
                    "description": "VizieR 表名，describe/query/cone/export 时使用，如 'I/355/gaiadr3'"
                },
                "adql": {
                    "type": "string",
                    "description": "query/export 时的 ADQL 语句（与 table 二选一）"
                },
                "columns": {
                    "type": "string",
                    "description": "query/export + table 模式下指定列（逗号分隔，默认 *）"
                },
                "coords": {
                    "type": "object",
                    "description": "cone 时的坐标参数",
                    "properties": {
                        "ra": { "type": "number", "description": "赤经 RA（度）" },
                        "dec": { "type": "number", "description": "赤纬 Dec（度）" },
                        "radius_deg": { "type": "number", "description": "检索半径（度），默认 0.1", "default": 0.1 },
                        "strategy": { "type": "string", "enum": ["nearest", "all"], "description": "nearest（默认）/ all", "default": "nearest" }
                    },
                    "required": ["ra", "dec"]
                },
                "limit": {
                    "type": "integer",
                    "description": "最大返回行数（默认 50，上限 2000）",
                    "default": 50
                },
                "bibcode": {
                    "type": "string",
                    "description": "lookup 时的 ADS bibcode"
                }
            },
            "required": ["action"]
        })
    }

    fn group(&self) -> &str {
        "as:research"
    }

    fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
        true
    }

    fn is_readonly(&self) -> bool {
        false
    }

    async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
        let action = match args.get("action").and_then(|v| v.as_str()) {
            Some(a) => a,
            None => return ToolOutput::error("缺少必需参数 'action'"),
        };

        let state = &ctx.app_state;

        match action {
            "search" => do_catalog_search(state, &args).await,
            "describe" => do_catalog_describe(state, &args).await,
            "query" => do_catalog_query(state, &args).await,
            "cone" => do_catalog_cone(state, &args).await,
            "export" => do_catalog_export(state, &args).await,
            "lookup" => do_catalog_lookup(state, &args).await,
            other => ToolOutput::error(format!(
                "未知 action '{}'，可选: search/describe/query/cone/export/lookup",
                other
            )),
        }
    }
}

async fn do_catalog_search(state: &crate::api::AppState, args: &serde_json::Value) -> ToolOutput {
    let keyword = match args.get("keyword").and_then(|v| v.as_str()) {
        Some(k) => k,
        None => return ToolOutput::error("search 需要 'keyword' 参数"),
    };
    let limit = args
        .get("limit")
        .and_then(|v| v.as_i64())
        .unwrap_or(10)
        .clamp(1, 30) as usize;

    let catalog =
        crate::services::cds::vizier::VizierCatalog::new(&state.db, &state.sources.vizier);
    let results = match catalog.search(keyword, limit).await {
        Ok(r) => r,
        Err(e) => return ToolOutput::error(format!("搜索失败: {}", e)),
    };

    if results.is_empty() {
        return ToolOutput::success(format!("未找到与 '{}' 匹配的星表", keyword), json!([]));
    }

    let mut content = format!(
        "搜索 '{}' 匹配到 {} 个星表（按数据量降序）：\n\n",
        keyword,
        results.len()
    );
    let items: Vec<serde_json::Value> = results
        .iter()
        .map(|entry| {
            let nrows_str = entry.nrows.map(|n| format!("{} 行", n)).unwrap_or_else(|| "行数未知".into());
            content.push_str(&format!("- `{}` — {} ({})\n", entry.table_name, entry.description, nrows_str));
            json!({"table": entry.table_name, "description": entry.description, "nrows": entry.nrows})
        })
        .collect();

    ToolOutput::success(content, json!(items))
}

async fn do_catalog_describe(state: &crate::api::AppState, args: &serde_json::Value) -> ToolOutput {
    let table = match args.get("table").and_then(|v| v.as_str()) {
        Some(t) => t,
        None => return ToolOutput::error("describe 需要 'table' 参数"),
    };

    let catalog =
        crate::services::cds::vizier::VizierCatalog::new(&state.db, &state.sources.vizier);
    let columns = match catalog.describe(table).await {
        Ok(c) => c,
        Err(e) => return ToolOutput::error(format!("查询表结构失败: {}", e)),
    };

    if columns.is_empty() {
        return ToolOutput::success(format!("表 '{}' 未找到列定义", table), json!([]));
    }

    let mut content = format!("表 `{}` 共 {} 列：\n\n", table, columns.len());
    let items: Vec<serde_json::Value> = columns
        .iter()
        .map(|col| {
            let unit_str = col.unit.as_deref().unwrap_or("");
            let desc_str = col.description.as_deref().unwrap_or("");
            content.push_str(&format!("- `{}` ({}) {} — {}\n", col.column_name, col.datatype, unit_str, desc_str));
            json!({"name": col.column_name, "datatype": col.datatype, "unit": col.unit, "description": col.description})
        })
        .collect();

    ToolOutput::success(content, json!(items))
}

async fn do_catalog_query(state: &crate::api::AppState, args: &serde_json::Value) -> ToolOutput {
    let limit = args
        .get("limit")
        .and_then(|v| v.as_i64())
        .unwrap_or(50)
        .clamp(1, 2000);

    let result = if let Some(adql) = args.get("adql").and_then(|v| v.as_str()) {
        info!(
            "[CatalogOp:query] ADQL (limit={}): {}",
            limit,
            adql.chars().take(150).collect::<String>()
        );
        crate::services::cds::vizier::query_adql_cached(
            &state.db,
            &state.sources.vizier,
            adql,
            limit,
        )
        .await
    } else if let Some(table) = args.get("table").and_then(|v| v.as_str()) {
        let columns: Vec<String> = args
            .get("columns")
            .and_then(|v| v.as_str())
            .map(|c| c.split(',').map(|s| s.trim().to_string()).collect())
            .unwrap_or_default();
        info!("[CatalogOp:query] table={} limit={}", table, limit);
        crate::services::cds::vizier::query_table(
            &state.db,
            &state.sources.vizier,
            table,
            &columns,
            limit,
        )
        .await
    } else {
        return ToolOutput::error("query 需要 'adql' 或 'table' 参数之一");
    };

    match result {
        Ok(r) => {
            let content = render_result_table(&r, 20);
            ToolOutput::success(content, json!(r))
        }
        Err(e) => ToolOutput::error(format!("查询失败: {}", e)),
    }
}

async fn do_catalog_cone(state: &crate::api::AppState, args: &serde_json::Value) -> ToolOutput {
    let coords = match args.get("coords").and_then(|v| v.as_object()) {
        Some(c) => c,
        None => return ToolOutput::error("cone 需要 'coords' 参数（含 ra/dec）"),
    };
    let ra = match coords.get("ra").and_then(|v| v.as_f64()) {
        Some(v) => v,
        None => return ToolOutput::error("coords 缺少 'ra'"),
    };
    let dec = match coords.get("dec").and_then(|v| v.as_f64()) {
        Some(v) => v,
        None => return ToolOutput::error("coords 缺少 'dec'"),
    };
    let radius = coords
        .get("radius_deg")
        .and_then(|v| v.as_f64())
        .unwrap_or(0.1);
    let nearest = coords.get("strategy").and_then(|v| v.as_str()) != Some("all");
    let table = match args.get("table").and_then(|v| v.as_str()) {
        Some(t) => t,
        None => return ToolOutput::error("cone 需要 'table' 参数"),
    };
    let max_records = args
        .get("limit")
        .and_then(|v| v.as_i64())
        .unwrap_or(50)
        .clamp(1, 2000);

    info!(
        "[CatalogOp:cone] ra={} dec={} radius={}° table={} strategy={}",
        ra,
        dec,
        radius,
        table,
        if nearest { "nearest" } else { "all" }
    );

    match crate::services::cds::vizier::cone_search(
        &state.db,
        &state.sources.vizier,
        ra,
        dec,
        radius,
        table,
        max_records,
        nearest,
    )
    .await
    {
        Ok(r) => {
            let mut content = format!(
                "Cone Search 结果 (ra={}, dec={}, radius={}°)\n\n",
                ra, dec, radius
            );
            content.push_str(&render_result_table(&r, 20));
            ToolOutput::success(content, json!(r))
        }
        Err(e) => ToolOutput::error(format!("Cone Search 失败: {}", e)),
    }
}

async fn do_catalog_export(state: &crate::api::AppState, args: &serde_json::Value) -> ToolOutput {
    let limit = args
        .get("limit")
        .and_then(|v| v.as_i64())
        .unwrap_or(100)
        .clamp(1, 5000);

    let adql = args.get("adql").and_then(|v| v.as_str());
    let table = args.get("table").and_then(|v| v.as_str());
    let columns = args.get("columns").and_then(|v| v.as_str());

    if adql.is_none() && table.is_none() {
        return ToolOutput::error("export 需要 'adql' 或 'table' 参数");
    }

    let catalog =
        crate::services::cds::vizier::VizierCatalog::new(&state.db, &state.sources.vizier);
    let result = match catalog
        .export_to_file(
            state.config.storage.library_dir.to_str().unwrap_or("."),
            adql,
            table,
            columns,
            limit,
        )
        .await
    {
        Ok(r) => r,
        Err(e) => return ToolOutput::error(format!("导出失败: {}", e)),
    };

    ToolOutput::success(
        format!(
            "已导出 {} 行数据到 `{}`\n文件大小: {}",
            result.row_count,
            result.path.display(),
            result.size_bytes,
        ),
        json!({
            "path": result.path.to_str(),
            "rows": result.row_count,
            "columns": result.column_count,
        }),
    )
}

async fn do_catalog_lookup(state: &crate::api::AppState, args: &serde_json::Value) -> ToolOutput {
    let bibcode = match args.get("bibcode").and_then(|v| v.as_str()) {
        Some(b) => b,
        None => return ToolOutput::error("lookup 需要 'bibcode' 参数"),
    };
    let limit = args
        .get("limit")
        .and_then(|v| v.as_i64())
        .unwrap_or(10)
        .clamp(1, 30) as usize;

    let catalog = crate::services::cds::vizier::VizierCatalog::with_ads(
        &state.db,
        &state.sources.vizier,
        &state.sources.ads,
    );
    let results = match catalog.lookup(bibcode, limit).await {
        Ok(r) => r,
        Err(e) => return ToolOutput::error(format!("查找失败: {}", e)),
    };

    if results.is_empty() {
        return ToolOutput::success(
            format!("文献 '{}' 未在 CDS/VizieR 中找到关联数据表", bibcode),
            json!([]),
        );
    }

    let mut content = format!(
        "文献 '{}' 关联 {} 个 VizieR 数据表：\n\n",
        bibcode,
        results.len()
    );
    let items: Vec<serde_json::Value> = results
        .iter()
        .map(|entry| {
            let nrows_str = entry.nrows.map(|n| format!("{} 行", n)).unwrap_or_default();
            content.push_str(&format!("- `{}` — {} ({})\n", entry.table_name, entry.description, nrows_str));
            json!({"table": entry.table_name, "description": entry.description, "nrows": entry.nrows})
        })
        .collect();

    ToolOutput::success(content, json!(items))
}