AstroResearch/src/services/observation/gaia.rs

// src/services/observation/gaia.rs
//
// Gaia ObservationFetcher 实现
//
// 注册的 fetcher：
//   GaiaSpectrumFetcher   —— (Gaia, Spectrum)，XP_CONTINUOUS/XP_SAMPLED/RVS
//   GaiaLightCurveFetcher —— (Gaia, LightCurve)，EPOCH_PHOTOMETRY（G/BP/RP 三波段多文件）
//
// 关键修复：EPOCH_PHOTOMETRY 的 ZIP 内含 G/BP/RP 三个独立文件，
// extract_all_zip_entries 返回全部而非首个（旧 extract_first_zip_entry 丢失 BP/RP 波段）。

use crate::api::AppState;
use crate::clients::gaia::{GaiaRetrievalType, GaiaSourceRow};
use crate::services::observation::cache::{
    cached_files_total_size, fetch_observation_cache, file_url_from_path, log_write_failure,
    persist_bytes, write_observation_cache,
};
use crate::services::observation::fetcher::{Candidate, ObservationFetcher};
use crate::services::observation::types::{
    Artifact, ObservationProduct, ProductSpec, ProductType, Source,
};
use anyhow::{anyhow, Result};
use async_trait::async_trait;
use std::io::{Cursor, Read};
use std::time::Duration;
use tracing::{info, warn};

// ── 版本与解析辅助 ──

#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum GaiaRelease { #[default] Dr3 }

impl GaiaRelease {
    pub fn valid_values() -> &'static [&'static str] {
        &["dr3"]
    }
    pub fn dr_str(&self) -> &'static str { "dr3" }
    pub fn tap_table(&self) -> &'static str { "gaiadr3.gaia_source" }
    pub fn display(&self) -> &'static str { "DR3" }
}

pub fn parse_gaia_release(s: Option<&str>) -> Result<GaiaRelease> {
    match s {
        None | Some("dr3") => Ok(GaiaRelease::Dr3),
        Some(other) => Err(anyhow!("不支持的 Gaia release '{}'，可选: {:?}", other, GaiaRelease::valid_values())),
    }
}

pub fn parse_gaia_retrieval_type_str(s: Option<&str>) -> Result<GaiaRetrievalType> {
    match s.map(|x| x.to_uppercase()).as_deref() {
        None => Ok(GaiaRetrievalType::XpContinuous),
        Some("XP_CONTINUOUS") => Ok(GaiaRetrievalType::XpContinuous),
        Some("XP_SAMPLED") => Ok(GaiaRetrievalType::XpSampled),
        Some("EPOCH_PHOTOMETRY") => Ok(GaiaRetrievalType::EpochPhotometry),
        Some("RVS") => Ok(GaiaRetrievalType::Rvs),
        Some(other) => Err(anyhow!("不支持的 Gaia subtype '{}'，可选: {:?}", other, GaiaRetrievalType::valid_values())),
    }
}

pub fn parse_gaia_retrieval_type(s: &str) -> Result<GaiaRetrievalType> {
    use GaiaRetrievalType::*;
    Ok(match s.trim().to_uppercase().as_str() {
        "XP_CONTINUOUS" => XpContinuous,
        "XP_SAMPLED" => XpSampled,
        "EPOCH_PHOTOMETRY" => EpochPhotometry,
        "RVS" => Rvs,
        other => return Err(anyhow!("不支持的 Gaia retrieval_type '{}'，可选: {:?}", other, GaiaRetrievalType::valid_values())),
    })
}

fn validate_cone_params(ra: f64, dec: f64, radius_deg: f64) -> Result<()> {
    // Gaia TAP/ADQL 无硬性半径上限，但同步查询限 2000 行、大半径易超时
    if !(radius_deg > 0.0 && radius_deg <= 30.0) {
        return Err(anyhow!(
            "检索半径应在 0~30 度之间（不含 0），当前: {}。Gaia 主表很大，建议 ≤1° 以内避免同步查询超时",
            radius_deg
        ));
    }
    if radius_deg > 1.0 {
        warn!(
            "[Observation] Gaia cone radius {}° 超出建议范围（≤1°），gaiadr3 主表很大易导致同步查询超时",
            radius_deg
        );
    }
    if !(-360.0..=360.0).contains(&ra) || !(-90.0..=90.0).contains(&dec) {
        return Err(anyhow!("坐标范围异常 (ra={}, dec={})", ra, dec));
    }
    Ok(())
}

fn row_to_candidate(row: &GaiaSourceRow) -> Candidate {
    Candidate {
        source: Source::Gaia,
        source_id: row.source_id.clone(),
        label: format!("Gaia {}", row.source_id),
        ra: row.ra,
        dec: row.dec,
        distance: row.distance,
        raw: Some(serde_json::to_value(row).unwrap_or(serde_json::Value::Null)),
    }
}

/// 通用 Gaia 下载逻辑（spectrum 与 lightcurve 共用，仅 retrieval_type 与产品类型不同）
async fn download_gaia_product(
    state: &AppState,
    source_id: &str,
    retrieval_type: GaiaRetrievalType,
    product_type: ProductType,
    force: bool,
) -> Result<ObservationProduct> {
    if source_id.trim().is_empty() {
        return Err(anyhow!("source_id 不能为空"));
    }
    let rt_param = retrieval_type.as_param();
    let rt_lower = rt_param.to_lowercase();
    let product = ProductSpec::with_subtype(product_type, rt_lower.clone());
    let dr = "dr3";
    let cache_key = format!("{}|{}|{}", dr, rt_param, source_id);
    let source_label = format!("Gaia {} ({})", source_id, rt_param);

    // 1) 缓存命中
    if !force {
        if let Some((artifacts, meta)) = fetch_observation_cache(&state.db, Source::Gaia, &product, &cache_key).await? {
            if cached_files_total_size(&state.config.library_dir, &artifacts).is_some() {
                info!("[Gaia] 缓存命中 (source_id={}, type={})", source_id, rt_param);
                let artifacts = artifacts.into_iter().map(|mut a| {
                    if a.size_bytes == 0 {
                        if let Ok(m) = std::fs::metadata(state.config.library_dir.join(&a.file_path)) {
                            a.size_bytes = m.len() as usize;
                        }
                    }
                    a
                }).collect();
                return Ok(ObservationProduct {
                    source: Source::Gaia, product,
                    source_id: cache_key, source_label, artifacts, source_meta: meta,
                });
            }
            warn!("[Gaia] 缓存记录存在但文件缺失，重新下载 (source_id={})", source_id);
        }
    }

    // 2) DataLink 下载 ZIP + 解包取**全部**文件（修复 EPOCH_PHOTOMETRY 多文件丢失）
    tokio::time::sleep(Duration::from_millis(50)).await;
    let zip_bytes = state.gaia.download_products(&[source_id.to_string()], retrieval_type, "fits").await?;
    let extracted = extract_all_zip_entries(&zip_bytes, source_id)?;

    // 3) 落盘
    // 光谱：单文件 → Telescope/gaia/spectrum/{rt}/{source_id}.{ext}
    // 光变：多文件 → Telescope/gaia/lightcurve/{rt}/{source_id}/{orig_name}
    let mut artifacts = Vec::with_capacity(extracted.len());
    for ext_file in extracted {
        let rel_path = if product_type == ProductType::LightCurve {
            let safe_name = ext_file.name.replace('/', "_");
            format!("Telescope/gaia/lightcurve/{rt}/{sid}/{name}",
                rt = rt_lower, sid = source_id, name = safe_name)
        } else {
            format!("Telescope/gaia/spectrum/{rt}/{sid}.{ext}",
                rt = rt_lower, sid = source_id, ext = ext_file.ext)
        };
        let size = persist_bytes(&state.config.library_dir, &rel_path, &ext_file.bytes)?;
        info!("[Gaia] 已保存 {} ({}B) → {}", ext_file.name, size, rel_path);
        artifacts.push(Artifact {
            band: parse_band_from_name(&ext_file.name),
            original_name: Some(ext_file.name.clone()),
            file_path: rel_path,
            file_url: String::new(),
            file_format: ext_file.ext.to_string(),
            size_bytes: size,
            cached: false,
        });
    }
    for a in &mut artifacts {
        a.file_url = file_url_from_path(&a.file_path);
    }
    if artifacts.is_empty() {
        return Err(anyhow!("Gaia DataLink ZIP 内无可下载文件条目 (source_id={}, type={})", source_id, rt_param));
    }

    let meta = serde_json::json!({"source_id": source_id, "retrieval_type": rt_param, "dr": dr});
    let meta_str = meta.to_string();
    if let Err(e) = write_observation_cache(
        &state.db, Source::Gaia, &product, &cache_key, None, None,
        &artifacts, Some(&meta_str),
    ).await {
        log_write_failure(Source::Gaia.as_str(), product.product.as_str(), e);
    }

    Ok(ObservationProduct {
        source: Source::Gaia, product,
        source_id: cache_key, source_label,
        artifacts, source_meta: Some(meta),
    })
}

// ═══════════════════════════════════════════════════════════════
// ZIP 解包（全量提取，修复历史 EPOCH_PHOTOMETRY 多文件丢失 bug）
// ═══════════════════════════════════════════════════════════════

struct ExtractedFile {
    name: String,
    ext: &'static str,
    bytes: Vec<u8>,
}

/// 从 ZIP 字节流中提取**全部**非目录文件条目
///
/// 旧实现 `extract_first_zip_entry` 只取首个文件，丢失了 BP/RP 波段——本函数修复之。
fn extract_all_zip_entries(zip_bytes: &[u8], source_id: &str) -> Result<Vec<ExtractedFile>> {
    let cursor = Cursor::new(zip_bytes);
    let mut archive = zip::ZipArchive::new(cursor).map_err(|e| anyhow!("Gaia ZIP 解析失败: {}", e))?;

    let mut out = Vec::new();
    for i in 0..archive.len() {
        let mut file = archive.by_index(i)
            .map_err(|e| anyhow!("读取 Gaia ZIP 条目 {} 失败: {}", i, e))?;
        if file.is_dir() { continue; }
        let name = file.name().to_string();
        let mut buf = Vec::new();
        file.read_to_end(&mut buf)
            .map_err(|e| anyhow!("解压 Gaia ZIP 条目失败: {}", e))?;
        if buf.is_empty() { continue; }

        let lower = name.to_lowercase();
        let ext = if lower.ends_with(".fits") { "fits" }
            else if lower.ends_with(".vot") || lower.ends_with(".xml") { "vot" }
            else if lower.ends_with(".csv") { "csv" }
            else { "fits" };

        out.push(ExtractedFile { name, ext, bytes: buf });
    }

    if out.is_empty() {
        return Err(anyhow!("Gaia ZIP 内无可下载文件条目 (source_id={})", source_id));
    }
    Ok(out)
}

/// 从 Gaia DataLink 文件名中提取波段标签
fn parse_band_from_name(name: &str) -> Option<String> {
    let upper = name.to_uppercase();
    for band in ["G", "BP", "RP"] {
        if upper.contains(&format!("_{}.", band)) || upper.contains(&format!("_{}_", band)) {
            return Some(band.to_string());
        }
    }
    None
}

// ═══════════════════════════════════════════════════════════════
// GaiaSpectrumFetcher
// ═══════════════════════════════════════════════════════════════

#[derive(Debug)]
pub struct GaiaSpectrumFetcher;

#[async_trait]
impl ObservationFetcher for GaiaSpectrumFetcher {
    fn key(&self) -> (Source, ProductType) {
        (Source::Gaia, ProductType::Spectrum)
    }

    fn subtypes(&self) -> &'static [&'static str] {
        &["xp_continuous", "xp_sampled", "rvs"]
    }

    fn releases(&self) -> &'static [&'static str] {
        GaiaRelease::valid_values()
    }

    fn default_release(&self) -> Option<&'static str> {
        Some("dr3")
    }

    fn suggested_max_radius_deg(&self) -> f64 {
        // 与 validate_cone_params 的建议范围一致
        1.0
    }

    fn identifier_format(&self) -> Option<&'static str> {
        Some("'XP_CONTINUOUS|source_id'（可省略 RT 前缀，默认 XP_CONTINUOUS）")
    }

    async fn cone_search_raw(
        &self,
        state: &AppState,
        ra: f64, dec: f64, radius_deg: f64,
        _release: Option<&str>, _subtype: Option<&str>,
    ) -> Result<Vec<Candidate>> {
        validate_cone_params(ra, dec, radius_deg)?;
        let rel = parse_gaia_release(_release)?;
        let result = state.gaia.cone_search(ra, dec, radius_deg, 50, true, rel).await?;
        Ok(result.rows.iter().map(row_to_candidate).collect())
    }

    async fn resolve_identifier(
        &self,
        identifier: &str,
        _release: Option<&str>,
        subtype: Option<&str>,
    ) -> Result<Candidate> {
        let (rt_param, gaia_sid) = match identifier.split_once('|') {
            Some((rt, s)) => (rt, s),
            None => ("XP_CONTINUOUS", identifier),
        };
        let rt = if subtype.is_some() {
            parse_gaia_retrieval_type_str(subtype)?
        } else {
            parse_gaia_retrieval_type(rt_param)?
        };
        if gaia_sid.trim().is_empty() {
            return Err(anyhow!("source_id 不能为空"));
        }
        Ok(Candidate {
            source: Source::Gaia,
            source_id: gaia_sid.to_string(),
            label: format!("Gaia {}", gaia_sid),
            ra: None, dec: None, distance: None,
            raw: Some(serde_json::json!({"source_id": gaia_sid, "retrieval_type": rt.as_param()})),
        })
    }

    async fn fetch(
        &self,
        state: &AppState,
        candidate: &Candidate,
        _release: Option<&str>,
        subtype: Option<&str>,
        force: bool,
    ) -> Result<ObservationProduct> {
        let source_id = candidate.raw.as_ref()
            .and_then(|v| v.get("source_id"))
            .and_then(|v| v.as_str())
            .unwrap_or(&candidate.source_id);
        let rt = if let Some(rt_str) = candidate.raw.as_ref()
            .and_then(|v| v.get("retrieval_type"))
            .and_then(|v| v.as_str())
        {
            parse_gaia_retrieval_type(rt_str)?
        } else {
            parse_gaia_retrieval_type_str(subtype)?
        };
        download_gaia_product(state, source_id, rt, ProductType::Spectrum, force).await
    }
}

// ═══════════════════════════════════════════════════════════════
// GaiaLightCurveFetcher（EPOCH_PHOTOMETRY，多文件产物）
// ═══════════════════════════════════════════════════════════════

#[derive(Debug)]
pub struct GaiaLightCurveFetcher;

#[async_trait]
impl ObservationFetcher for GaiaLightCurveFetcher {
    fn key(&self) -> (Source, ProductType) {
        (Source::Gaia, ProductType::LightCurve)
    }

    fn subtypes(&self) -> &'static [&'static str] {
        &["epoch_photometry"]
    }

    fn releases(&self) -> &'static [&'static str] {
        GaiaRelease::valid_values()
    }

    fn default_release(&self) -> Option<&'static str> {
        Some("dr3")
    }

    fn suggested_max_radius_deg(&self) -> f64 {
        // 与 validate_cone_params 的建议范围一致
        1.0
    }

    fn identifier_format(&self) -> Option<&'static str> {
        Some("source_id，如 '65214031805717376'")
    }

    async fn cone_search_raw(
        &self,
        state: &AppState,
        ra: f64, dec: f64, radius_deg: f64,
        _release: Option<&str>, _subtype: Option<&str>,
    ) -> Result<Vec<Candidate>> {
        // 光变 cone 复用 Gaia 主表检索（不强制 xp_only，因为光变对所有源都有）
        validate_cone_params(ra, dec, radius_deg)?;
        let rel = parse_gaia_release(_release)?;
        let result = state.gaia.cone_search(ra, dec, radius_deg, 50, false, rel).await?;
        Ok(result.rows.iter().map(row_to_candidate).collect())
    }

    async fn resolve_identifier(
        &self,
        identifier: &str,
        _release: Option<&str>,
        _subtype: Option<&str>,
    ) -> Result<Candidate> {
        let sid = identifier.trim();
        if sid.is_empty() {
            return Err(anyhow!("source_id 不能为空"));
        }
        Ok(Candidate {
            source: Source::Gaia,
            source_id: sid.to_string(),
            label: format!("Gaia {}", sid),
            ra: None, dec: None, distance: None,
            raw: Some(serde_json::json!({"source_id": sid, "retrieval_type": "EPOCH_PHOTOMETRY"})),
        })
    }

    async fn fetch(
        &self,
        state: &AppState,
        candidate: &Candidate,
        _release: Option<&str>,
        _subtype: Option<&str>,
        force: bool,
    ) -> Result<ObservationProduct> {
        let source_id = candidate.raw.as_ref()
            .and_then(|v| v.get("source_id"))
            .and_then(|v| v.as_str())
            .unwrap_or(&candidate.source_id);
        download_gaia_product(state, source_id, GaiaRetrievalType::EpochPhotometry, ProductType::LightCurve, force).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;

    #[test]
    fn test_validate_cone_params() {
        // 合法范围：0~30°（建议 ≤1°，超出仅警告）
        assert!(validate_cone_params(180.0, 30.0, 0.1).is_ok());
        assert!(validate_cone_params(180.0, 30.0, 1.0).is_ok()); // 建议上限内
        assert!(validate_cone_params(180.0, 30.0, 5.0).is_ok()); // 超建议但仍合法
        assert!(validate_cone_params(180.0, 30.0, 31.0).is_err()); // 超硬性上限
    }

    #[test]
    fn test_parse_band_from_name_epoch() {
        assert_eq!(parse_band_from_name("65214031805717376_EPOCH_PHOTOMETRY_G.fits").as_deref(), Some("G"));
        assert_eq!(parse_band_from_name("65214031805717376_EPOCH_PHOTOMETRY_BP.fits").as_deref(), Some("BP"));
        assert_eq!(parse_band_from_name("65214031805717376_EPOCH_PHOTOMETRY_RP.fits").as_deref(), Some("RP"));
    }

    #[test]
    fn test_parse_band_from_name_xp() {
        assert_eq!(parse_band_from_name("65214031805717376_XP_CONTINUOUS.fits"), None);
    }

    #[test]
    fn test_extract_all_zip_entries_single() {
        let mut buf = Vec::new();
        {
            let mut writer = zip::ZipWriter::new(std::io::Cursor::new(&mut buf));
            let opts = zip::write::SimpleFileOptions::default();
            writer.start_file("123_XP_CONTINUOUS.fits", opts).unwrap();
            writer.write_all(b"SIMPLE").unwrap();
            writer.finish().unwrap();
        }
        let files = extract_all_zip_entries(&buf, "123").unwrap();
        assert_eq!(files.len(), 1);
        assert_eq!(files[0].ext, "fits");
    }

    #[test]
    fn test_extract_all_zip_entries_multi_band() {
        // 历史 bug 验证：EPOCH_PHOTOMETRY 三波段，旧实现只取首个 G 波段
        let mut buf = Vec::new();
        {
            let mut writer = zip::ZipWriter::new(std::io::Cursor::new(&mut buf));
            let opts = zip::write::SimpleFileOptions::default();
            writer.start_file("123_EPOCH_PHOTOMETRY_G.fits", opts).unwrap();
            writer.write_all(b"G_DATA").unwrap();
            writer.start_file("123_EPOCH_PHOTOMETRY_BP.fits", opts).unwrap();
            writer.write_all(b"BP_DATA").unwrap();
            writer.start_file("123_EPOCH_PHOTOMETRY_RP.fits", opts).unwrap();
            writer.write_all(b"RP_DATA").unwrap();
            writer.finish().unwrap();
        }
        let files = extract_all_zip_entries(&buf, "123").unwrap();
        assert_eq!(files.len(), 3, "EPOCH_PHOTOMETRY 应返回全部 3 个波段文件");
        let bands: Vec<_> = files.iter().map(|f| parse_band_from_name(&f.name)).collect();
        assert!(bands.contains(&Some("G".into())));
        assert!(bands.contains(&Some("BP".into())));
        assert!(bands.contains(&Some("RP".into())));
    }
}