AstroResearch/src/services/observation/desi.rs

// src/services/observation/desi.rs
//
// DESI ObservationFetcher 实现
//
// 注册的 fetcher：
//   DesiSpectrumFetcher —— (Desi, Spectrum)，HEALPix coadd（一像素含多目标光谱）

use crate::api::AppState;
use crate::clients::desi::DesiSpectrumRow;
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::time::Duration;
use tracing::{info, warn};

// ── 版本枚举 ──

#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum DesiRelease {
    Edr,
    #[default]
    Dr1,
}

impl DesiRelease {
    pub fn valid_values() -> &'static [&'static str] {
        &["edr", "dr1"]
    }
    pub fn dr_segment(&self) -> &'static str {
        match self {
            Self::Edr => "edr",
            Self::Dr1 => "dr1",
        }
    }
    pub fn specredux_ver(&self) -> &'static str {
        match self {
            Self::Edr => "fuji",
            Self::Dr1 => "iron",
        }
    }
    pub fn datalab_table(&self) -> &'static str {
        match self {
            Self::Edr => "desi_edr.zpix",
            Self::Dr1 => "desi_dr1.zpix",
        }
    }
    pub fn display(&self) -> &'static str {
        match self {
            Self::Edr => "EDR",
            Self::Dr1 => "DR1",
        }
    }
}

pub fn parse_desi_release(s: Option<&str>) -> Result<DesiRelease> {
    Ok(match s.map(|x| x.to_lowercase()).as_deref() {
        None => DesiRelease::Dr1,
        Some("edr") => DesiRelease::Edr,
        Some("dr1") => DesiRelease::Dr1,
        Some(other) => {
            return Err(anyhow!(
                "不支持的 DESI release '{}'，可选: {:?}",
                other,
                DesiRelease::valid_values()
            ))
        }
    })
}

fn validate_cone_params(ra: f64, dec: f64, radius_deg: f64) -> Result<()> {
    // NOIRLab Data Lab 无硬性半径上限，但大半径易超时
    if !(radius_deg > 0.0 && radius_deg <= 30.0) {
        return Err(anyhow!(
            "检索半径应在 0~30 度之间（不含 0），当前: {}。DESI zpix 建议检索半径 ≤1° 以避免超时",
            radius_deg
        ));
    }
    if radius_deg > 1.0 {
        warn!(
            "[Observation] DESI cone radius {}° 超出建议范围（≤1°），大半径查询可能被服务端超时拒绝",
            radius_deg
        );
    }
    if !(-360.0..=360.0).contains(&ra) || !(-90.0..=90.0).contains(&dec) {
        return Err(anyhow!("坐标范围异常 (ra={}, dec={})", ra, dec));
    }
    Ok(())
}

// ═══════════════════════════════════════════════════════════════
// DesiSpectrumFetcher
// ═══════════════════════════════════════════════════════════════

#[derive(Debug)]
pub struct DesiSpectrumFetcher;

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

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

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

    fn default_release(&self) -> Option<&'static str> {
        // 与 parse_desi_release(None) 的默认值保持一致
        Some("dr1")
    }

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

    fn identifier_format(&self) -> Option<&'static str> {
        Some("'survey-program-healpix'，如 'main-dark-10050'")
    }

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

    async fn resolve_identifier(
        &self,
        identifier: &str,
        _release: Option<&str>,
        _subtype: Option<&str>,
        _version: Option<&str>,
    ) -> Result<Candidate> {
        // 格式 "{survey}-{program}-{healpix}"
        let parts: Vec<&str> = identifier.split('-').collect();
        if parts.len() != 3 {
            return Err(anyhow!(
                "DESI source_id 应为 'survey-program-healpix' 格式（如 'main-dark-10050'），得到 '{}'",
                identifier
            ));
        }
        let survey = parts[0];
        let program = parts[1];
        // 白名单校验：survey/program 会拼进文件系统路径和下载 URL，
        // 必须防止路径穿越（如 "../etc"）。仅允许 DESI 实际使用的取值。
        const VALID_SURVEYS: &[&str] = &["main", "sv1", "sv2", "sv3", "sv", "backup", "cumulative"];
        const VALID_PROGRAMS: &[&str] = &["dark", "bright", "backup", "other"];
        if !VALID_SURVEYS.contains(&survey) {
            return Err(anyhow!(
                "DESI survey '{}' 不在允许列表 {:?} 中",
                survey,
                VALID_SURVEYS
            ));
        }
        if !VALID_PROGRAMS.contains(&program) {
            return Err(anyhow!(
                "DESI program '{}' 不在允许列表 {:?} 中",
                program,
                VALID_PROGRAMS
            ));
        }
        let healpix: i64 = parts[2]
            .parse()
            .map_err(|_| anyhow!("healpix 非数字: '{}'", parts[2]))?;
        if healpix < 0 {
            return Err(anyhow!("healpix 应为非负整数，当前: {}", healpix));
        }
        Ok(Candidate {
            source: Source::Desi,
            source_id: identifier.to_string(),
            label: format!("{} {} healpix {}", survey, program, healpix),
            ra: None,
            dec: None,
            distance: None,
            raw: Some(serde_json::json!({
                "survey": survey, "program": program, "healpix": healpix
            })),
        })
    }

    async fn fetch(
        &self,
        state: &AppState,
        candidate: &Candidate,
        release: Option<&str>,
        _subtype: Option<&str>,
        _version: Option<&str>,
        force: bool,
    ) -> Result<ObservationProduct> {
        let survey = candidate
            .raw
            .as_ref()
            .and_then(|v| v.get("survey"))
            .and_then(|v| v.as_str())
            .ok_or_else(|| anyhow!("DESI Candidate 缺 survey"))?;
        let program = candidate
            .raw
            .as_ref()
            .and_then(|v| v.get("program"))
            .and_then(|v| v.as_str())
            .ok_or_else(|| anyhow!("DESI Candidate 缺 program"))?;
        let healpix = candidate
            .raw
            .as_ref()
            .and_then(|v| v.get("healpix"))
            .and_then(|v| v.as_i64())
            .ok_or_else(|| anyhow!("DESI Candidate 缺 healpix"))?;
        let rel = parse_desi_release(release)?;
        let product = ProductSpec::with_subtype(ProductType::Spectrum, "coadd");
        let dr = rel.dr_segment();
        let cache_key = format!("{}|{}-{}-{}", dr, survey, program, healpix);
        let source_label = format!("{} {} healpix {}", survey, program, healpix);

        // 1) 缓存命中
        if !force {
            if let Some((artifacts, meta)) =
                fetch_observation_cache(&state.db, Source::Desi, &product, &cache_key).await?
            {
                if cached_files_total_size(&state.config.library_dir, &artifacts)
                    .await
                    .is_some()
                {
                    info!("[DESI] coadd 缓存命中 ({})", cache_key);
                    return Ok(ObservationProduct {
                        source: Source::Desi,
                        product,
                        source_id: cache_key,
                        source_label,
                        artifacts,
                        source_meta: meta,
                    });
                }
                warn!("[DESI] 缓存记录存在但文件缺失，重新下载 ({})", cache_key);
            }
        }

        // 2) 下载 coadd FITS（无压缩）
        tokio::time::sleep(Duration::from_millis(50)).await;
        let fits_bytes = state
            .desi
            .download_coadd(survey, program, healpix, rel)
            .await?;
        let rel_path = format!(
            "Telescope/desi/{dr}/{survey}-{program}/{healpix}.fits",
            dr = dr,
            survey = survey,
            program = program,
            healpix = healpix
        );
        let size = persist_bytes(&state.config.library_dir, &rel_path, &fits_bytes).await?;
        info!(
            "[DESI] coadd 已保存 {} → {} ({}B)",
            cache_key, rel_path, size
        );

        let meta =
            serde_json::json!({"survey": survey, "program": program, "healpix": healpix, "dr": dr});
        let artifact = Artifact {
            band: None,
            original_name: None,
            file_path: rel_path.clone(),
            file_url: file_url_from_path(&rel_path),
            file_format: "fits".to_string(),
            size_bytes: size,
            cached: false,
        };
        let meta_str = meta.to_string();
        if let Err(e) = write_observation_cache(
            &state.db,
            Source::Desi,
            &product,
            &cache_key,
            None,
            None,
            std::slice::from_ref(&artifact),
            Some(&meta_str),
        )
        .await
        {
            log_write_failure(Source::Desi.as_str(), "spectrum", e);
        }

        Ok(ObservationProduct {
            source: Source::Desi,
            product,
            source_id: cache_key,
            source_label,
            artifacts: vec![artifact],
            source_meta: Some(meta),
        })
    }
}

fn row_to_candidate(row: &DesiSpectrumRow) -> Candidate {
    Candidate {
        source: Source::Desi,
        source_id: format!("{}-{}-{}", row.survey, row.program, row.healpix),
        label: format!("{} {} healpix {}", row.survey, row.program, row.healpix),
        ra: row.ra,
        dec: row.dec,
        distance: None,
        raw: Some(serde_json::json!({
            "survey": row.survey, "program": row.program, "healpix": row.healpix
        })),
    }
}

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

    #[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()); // 超硬性上限
        assert!(validate_cone_params(180.0, 91.0, 0.1).is_err()); // dec 越界
    }
}