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

// src/agent/tools/astro/research/observation.rs
//
// FindObservationTool —— 统一观测数据下载工具
//
// 跨 (LAMOST/Gaia/SDSS/DESI) × (Spectrum/LightCurve/Photometry/Image) 双轴：
//   - 坐标模式（默认）：给 ra/dec/radius + source + product + strategy，自动 cone 检索 → 选源 → 下载
//   - 标识符模式：给 source + product + source_ids，直接按标识下载（跳过检索）
//
// 源 × 产品支持矩阵（registry 注册决定，工具运行时按需校验）：
//              Spectrum    LightCurve   Photometry
//   LAMOST      lrs/mrs        -            -
//   Gaia        xp_continuous  epoch_photometry  -
//              xp_sampled
//              rvs
//   SDSS        spec/apstar     -            -
//              aspcap
//   DESI        coadd           -            -
//
// 源标识格式：
//   LAMOST spectrum: obsid 数字，如 "438809089"
//   Gaia spectrum: "XP_CONTINUOUS|source_id"（可省略 RT 前缀，默认 XP_CONTINUOUS）
//   Gaia lightcurve(epoch_photometry): source_id
//   SDSS spectrum(spec): "run2d-plate-mjd-fiberid"，如 "26-2225-53729-439"
//   SDSS spectrum(apstar/aspcap): "telescope|field|apogee_id"
//   DESI spectrum: "survey-program-healpix"，如 "main-dark-10050"
//
// 结果自动缓存（observation_cache，按 source+product+source_id 去重），重复查询不会重复下载。
// Gaia 光变(EPOCH_PHOTOMETRY) 返回 G/BP/RP 三波段各一个文件（多 artifact）。

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

use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
use crate::services::observation::{
    download_observation, FindStrategy, ObservationRequest, ProductSpec, ProductType, Source,
};

pub struct FindObservationTool;

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

    fn display_name(&self) -> &str {
        "观测数据获取"
    }

    fn description(&self) -> &str {
        "下载天文观测数据（LAMOST/Gaia/SDSS/DESI 的光谱/光变/测光/图像）。\n\
         坐标模式：给 ra/dec + source + product，自动 cone 检索并下载最近（或全部）命中源。\n\
         标识符模式：给 source + product + source_ids，直接按源标识下载。\n\
         结果自动缓存，重复查询不重复下载。"
    }

    fn parameters(&self) -> serde_json::Value {
        // enum 列表从枚举的 valid_values() 动态生成，新增源/产品时自动同步
        let source_enum: Vec<&str> = Source::valid_values().to_vec();
        let product_enum: Vec<&str> = ProductType::valid_values().to_vec();
        json!({
            "type": "object",
            "properties": {
                "source": {
                    "type": "string",
                    "enum": source_enum,
                    "description": "数据源"
                },
                "product": {
                    "type": "string",
                    "enum": product_enum,
                    "description": "产品类型，默认 spectrum",
                    "default": "spectrum"
                },
                "subtype": {
                    "type": "string",
                    "description": "产品子类型（可选）。各源可选值通过 GET /api/observation/capabilities 查询；LAMOST spectrum: lrs/mrs；Gaia spectrum: xp_continuous/xp_sampled/rvs；Gaia lightcurve: epoch_photometry；SDSS spectrum: spec/apstar/aspcap；DESI spectrum: coadd"
                },
                "ra": {
                    "type": "number",
                    "description": "赤经 RA（度，J2000/ICRS）。坐标模式必填"
                },
                "dec": {
                    "type": "number",
                    "description": "赤纬 Dec（度，J2000/ICRS）。坐标模式必填"
                },
                "radius_deg": {
                    "type": "number",
                    "description": "检索半径（度），坐标模式用，默认 0.1。建议值与硬上限可通过 GET /api/observation/capabilities 查询（LAMOST 建议≤5°，其他≤1°；硬上限 30°）",
                    "default": 0.1
                },
                "strategy": {
                    "type": "string",
                    "enum": ["nearest", "all"],
                    "description": "选源策略（坐标模式）：nearest（默认）/ all",
                    "default": "nearest"
                },
                "source_ids": {
                    "type": "array",
                    "items": { "type": "string" },
                    "description": "源标识列表（标识符模式）。提供时切换到标识符模式，忽略 ra/dec/radius/strategy"
                },
                "release": {
                    "type": "string",
                    "description": "数据发布版本（可选，留空用各源默认）。可选值与默认值通过 GET /api/observation/capabilities 查询"
                },
                "version": {
                    "type": "string",
                    "description": "数据发布的子版本（可选，仅 LAMOST 有意义）。留空用该 DR 的默认（最新公开）子版本。可选值通过 GET /api/observation/capabilities 的 release_versions 字段查询"
                },
                "force": {
                    "type": "boolean",
                    "description": "是否强制重新下载（忽略缓存），默认 false",
                    "default": false
                }
            },
            "required": ["source"]
        })
    }

    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 state = &ctx.app_state;

        let source = match args.get("source").and_then(|v| v.as_str()) {
            Some(s) => match Source::from_str(s) {
                Ok(src) => src,
                Err(e) => return ToolOutput::error(e),
            },
            None => return ToolOutput::error("缺少必需参数 'source'（lamost/gaia/sdss/desi）"),
        };
        let product_type = match args.get("product").and_then(|v| v.as_str()) {
            None => ProductType::Spectrum,
            Some(s) => match ProductType::from_str(s) {
                Ok(p) => p,
                Err(e) => return ToolOutput::error(e),
            },
        };
        let subtype = args
            .get("subtype")
            .and_then(|v| v.as_str())
            .map(|s| s.to_string());
        let product = ProductSpec {
            product: product_type,
            subtype,
        };
        let force = args.get("force").and_then(|v| v.as_bool()).unwrap_or(false);
        let release = args
            .get("release")
            .and_then(|v| v.as_str())
            .map(|s| s.to_string());
        let version = args
            .get("version")
            .and_then(|v| v.as_str())
            .map(|s| s.to_string());

        // 标识符模式 vs 坐标模式
        let request = if let Some(ids) = args.get("source_ids").and_then(|v| v.as_array()) {
            let identifiers: Vec<String> = ids
                .iter()
                .filter_map(|v| v.as_str().map(|s| s.to_string()))
                .collect();
            if identifiers.is_empty() {
                return ToolOutput::error("标识符模式下 source_ids 不能为空");
            }
            info!(
                "[FindObservation] by_id source={:?} product={:?} count={}",
                source,
                product.product,
                identifiers.len()
            );
            ObservationRequest::ByIdentifier {
                source,
                product,
                identifiers,
                release,
                version,
            }
        } else {
            let ra = match args.get("ra").and_then(|v| v.as_f64()) {
                Some(v) => v,
                None => {
                    return ToolOutput::error(
                        "坐标模式缺少必需参数 'ra'（赤经，度），或改用 source_ids 标识符模式",
                    )
                }
            };
            let dec = match args.get("dec").and_then(|v| v.as_f64()) {
                Some(v) => v,
                None => {
                    return ToolOutput::error(
                        "坐标模式缺少必需参数 'dec'（赤纬，度），或改用 source_ids 标识符模式",
                    )
                }
            };
            let radius = args
                .get("radius_deg")
                .and_then(|v| v.as_f64())
                .unwrap_or(0.1);
            let strategy = match args.get("strategy").and_then(|v| v.as_str()) {
                Some("all") => FindStrategy::All,
                _ => FindStrategy::Nearest,
            };
            info!(
                "[FindObservation] by_coords source={:?} product={:?} ra={} dec={} radius={}° strategy={:?}",
                source, product.product, ra, dec, radius, strategy
            );
            ObservationRequest::ByCoordinates {
                source,
                product,
                ra,
                dec,
                radius_deg: radius,
                strategy,
                release,
                version,
            }
        };

        match download_observation(state, &state.sources.observation_registry, &request, force)
            .await
        {
            Ok(batch) => {
                let content = render_batch(&batch);
                ToolOutput::success(content, json!(batch))
            }
            Err(e) => ToolOutput::error(format!("观测数据下载失败: {}", e)),
        }
    }
}

/// 渲染 ObservationBatch 为可读文本（支持多 artifact 展示）
fn render_batch(b: &crate::services::observation::ObservationBatch) -> String {
    let mut content = format!(
        "{} {} 下载",
        b.source.display(),
        b.product.product.display()
    );
    if let (Some(ra), Some(dec), Some(r)) = (b.ra, b.dec, b.radius_deg) {
        content.push_str(&format!("（ra={}, dec={}, radius={}°）", ra, dec, r));
    }
    if let Some(st) = &b.product.subtype {
        content.push_str(&format!(" [{}]", st));
    }
    content.push_str(&format!("：命中 {} 条\n", b.matched_count));

    if b.products.is_empty() && b.failures.is_empty() {
        content.push_str("(无观测数据覆盖或无匹配)\n");
        return content;
    }

    for p in &b.products {
        // 多 artifact 展示（如 Gaia 光变 G/BP/RP 三波段）
        let artifact_summary: Vec<String> = p
            .artifacts
            .iter()
            .map(|a| match &a.band {
                Some(band) => format!(
                    "{}波段 {} ({})",
                    band,
                    a.file_format.to_uppercase(),
                    format_size(a.size_bytes)
                ),
                None => format!(
                    "{} ({})",
                    a.file_format.to_uppercase(),
                    format_size(a.size_bytes)
                ),
            })
            .collect();
        content.push_str(&format!(
            "\n✓ 已下载（{}）: {} —— {}\n",
            if p.artifacts.iter().all(|a| a.cached) {
                "缓存"
            } else {
                "新下载"
            },
            p.source_label,
            artifact_summary.join(", ")
        ));
        for a in &p.artifacts {
            content.push_str(&format!(
                "  {}{}\n",
                a.file_path,
                a.band
                    .as_ref()
                    .map(|b| format!(" [{}]", b))
                    .unwrap_or_default()
            ));
        }
    }
    for f in &b.failures {
        content.push_str(&format!("\n✗ 下载失败 {}: {}\n", f.source_label, f.error));
    }
    content
}

fn format_size(bytes: usize) -> String {
    if bytes > 1024 * 1024 {
        format!("{:.1} MB", bytes as f64 / 1024.0 / 1024.0)
    } else if bytes > 1024 {
        format!("{:.1} KB", bytes as f64 / 1024.0)
    } else {
        format!("{} B", bytes)
    }
}