AstroResearch/src/services/observation/lamost.rs

// src/services/observation/lamost.rs
//
// LAMOST ObservationFetcher 实现
//
// 当前注册的 fetcher：
//   LamostSpectrumFetcher —— (Lamost, Spectrum)，LRS/MRS 双分辨率

use crate::api::AppState;
use crate::clients::lamost::LamostSpectrumRow;
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::Read;
use std::time::Duration;
use tracing::{info, warn};

// ── 版本与分辨率枚举（领域知识） ──

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LamostRelease {
    Dr5, Dr6, Dr7, Dr8, Dr9, Dr10, Dr11,
}

impl Default for LamostRelease {
    fn default() -> Self { Self::Dr10 }
}

impl LamostRelease {
    pub fn valid_values() -> &'static [&'static str] {
        &["dr5", "dr6", "dr7", "dr8", "dr9", "dr10", "dr11"]
    }
    pub fn path_segment(&self) -> &'static str {
        match self {
            Self::Dr5 => "dr5", Self::Dr6 => "dr6", Self::Dr7 => "dr7",
            Self::Dr8 => "dr8", Self::Dr9 => "dr9", Self::Dr10 => "dr10", Self::Dr11 => "dr11",
        }
    }
    pub fn version_segment(&self) -> &'static str {
        match self {
            Self::Dr5 | Self::Dr6 => "v2",
            Self::Dr7 => "v1.2",
            Self::Dr8 | Self::Dr9 | Self::Dr10 | Self::Dr11 => "v2.0",
        }
    }
    pub fn display(&self) -> &'static str {
        match self {
            Self::Dr5 => "DR5", Self::Dr6 => "DR6", Self::Dr7 => "DR7",
            Self::Dr8 => "DR8", Self::Dr9 => "DR9", Self::Dr10 => "DR10", Self::Dr11 => "DR11",
        }
    }
    pub fn supports_mrs(&self) -> bool {
        !matches!(self, Self::Dr5 | Self::Dr6)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LamostResolution { Lrs, Mrs }

impl LamostResolution {
    pub fn valid_values() -> &'static [&'static str] {
        &["lrs", "mrs"]
    }
    pub fn as_str(&self) -> &'static str {
        match self { Self::Lrs => "lrs", Self::Mrs => "mrs" }
    }
    pub fn display(&self) -> &'static str {
        match self { Self::Lrs => "LRS", Self::Mrs => "MRS" }
    }
}

pub fn parse_lamost_release(s: Option<&str>) -> Result<LamostRelease> {
    Ok(match s.map(|x| x.to_lowercase()).as_deref() {
        None => LamostRelease::Dr10,
        Some("dr5") => LamostRelease::Dr5,
        Some("dr6") => LamostRelease::Dr6,
        Some("dr7") => LamostRelease::Dr7,
        Some("dr8") => LamostRelease::Dr8,
        Some("dr9") => LamostRelease::Dr9,
        Some("dr10") => LamostRelease::Dr10,
        Some("dr11") => LamostRelease::Dr11,
        Some(other) => return Err(anyhow!("不支持的 LAMOST release '{}'，可选: {:?}", other, LamostRelease::valid_values())),
    })
}

pub fn parse_lamost_resolution(subtype: Option<&str>) -> Result<LamostResolution> {
    Ok(match subtype.map(|x| x.to_lowercase()).as_deref() {
        None => LamostResolution::Lrs,
        Some("lrs") | Some("low") => LamostResolution::Lrs,
        Some("mrs") | Some("medium") => LamostResolution::Mrs,
        Some(other) => return Err(anyhow!("不支持的 LAMOST subtype '{}'，可选: {:?}", other, LamostResolution::valid_values())),
    })
}

fn validate_cone_params(ra: f64, dec: f64, radius_deg: f64) -> Result<()> {
    if !(radius_deg > 0.0 && radius_deg <= 30.0) {
        return Err(anyhow!(
            "检索半径应在 0~30 度之间（不含 0），当前: {}。LAMOST 官方矩形检索建议 ≤10 平方度（≈半径 1.78°）",
            radius_deg
        ));
    }
    if radius_deg > 5.0 {
        warn!(
            "[Observation] LAMOST cone radius {}° 超出建议范围（≤5°），可能命中大量行导致响应缓慢",
            radius_deg
        );
    }
    if !(-360.0..=360.0).contains(&ra) || !(-90.0..=90.0).contains(&dec) {
        return Err(anyhow!("坐标范围异常 (ra={}, dec={})", ra, dec));
    }
    Ok(())
}

// ═══════════════════════════════════════════════════════════════
// LamostSpectrumFetcher
// ═══════════════════════════════════════════════════════════════

#[derive(Debug)]
pub struct LamostSpectrumFetcher;

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

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

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

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

    fn suggested_max_radius_deg(&self) -> f64 {
        // 与 validate_cone_params 的建议范围一致（≤5° 仅警告不拒绝）
        5.0
    }

    fn identifier_format(&self) -> Option<&'static str> {
        Some("obsid 数字，如 '438809089'")
    }

    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_lamost_release(release)?;
        let res = parse_lamost_resolution(subtype)?;
        if res == LamostResolution::Mrs && !rel.supports_mrs() {
            return Err(anyhow!("MRS（中分辨率）从 DR7 起才支持，{:?} 无 MRS 数据", rel));
        }
        let result = state.lamost.cone_search(ra, dec, radius_deg, rel, res).await?;
        Ok(result.rows.iter().map(|r| row_to_candidate(r, res.as_str())).collect())
    }

    async fn resolve_identifier(
        &self,
        identifier: &str,
        _release: Option<&str>,
        subtype: Option<&str>,
    ) -> Result<Candidate> {
        let obsid: i64 = identifier.trim().parse()
            .map_err(|_| anyhow!("LAMOST source_id 应为纯数字 obsid，得到 '{}'", identifier))?;
        if obsid <= 0 {
            return Err(anyhow!("obsid 应为正整数，当前: {}", obsid));
        }
        let res = parse_lamost_resolution(subtype)?.as_str().to_string();
        Ok(Candidate {
            source: Source::Lamost,
            source_id: obsid.to_string(),
            label: format!("obsid {}", obsid),
            ra: None,
            dec: None,
            distance: None,
            raw: Some(serde_json::json!({"obsid": obsid, "resolution": res})),
        })
    }

    async fn fetch(
        &self,
        state: &AppState,
        candidate: &Candidate,
        release: Option<&str>,
        subtype: Option<&str>,
        force: bool,
    ) -> Result<ObservationProduct> {
        let obsid = candidate.raw.as_ref()
            .and_then(|v| v.get("obsid"))
            .and_then(|v| v.as_i64())
            .ok_or_else(|| anyhow!("LAMOST Candidate 缺 obsid 字段"))?;
        let rel = parse_lamost_release(release)?;
        let res = parse_lamost_resolution(subtype)?;
        let product = ProductSpec::with_subtype(ProductType::Spectrum, res.as_str());
        let dr = rel.path_segment();
        let ver = rel.version_segment();
        let res_str = res.as_str();
        let cache_key = format!("{}|{}|{}|{}", dr, ver, res_str, obsid);
        let source_label = format!("obsid {}", obsid);

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

        // 2) 下载 + 解压
        tokio::time::sleep(Duration::from_millis(50)).await;
        let gz_bytes = state.lamost.download_fits(obsid, rel).await?;
        let fits_bytes = decompress_if_gzip(&gz_bytes)?;

        // 3) 落盘
        let rel_path = format!(
            "Telescope/lamost/spectrum/{res}/{dr}/{ver}/{obsid}.fits",
            res = res_str, dr = dr, ver = ver, obsid = obsid
        );
        let size = persist_bytes(&state.config.library_dir, &rel_path, &fits_bytes)?;
        info!("[LAMOST] FITS 已保存 obsid={} → {} ({}B)", obsid, rel_path, size);

        let meta = serde_json::json!({"obsid": obsid, "dr": dr, "version": ver, "resolution": res_str});
        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::Lamost, &product, &cache_key, None, None,
            std::slice::from_ref(&artifact), Some(&meta_str),
        ).await {
            log_write_failure(Source::Lamost.as_str(), "spectrum", e);
        }

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

fn row_to_candidate(row: &LamostSpectrumRow, resolution: &str) -> Candidate {
    let obsid = row.obsid;
    Candidate {
        source: Source::Lamost,
        source_id: obsid.to_string(),
        label: format!("obsid {}", obsid),
        ra: row.ra_obs,
        dec: row.dec_obs,
        distance: None,
        raw: Some(serde_json::json!({"obsid": obsid, "resolution": resolution})),
    }
}

fn decompress_if_gzip(bytes: &[u8]) -> Result<Vec<u8>> {
    if bytes.len() >= 2 && bytes[0] == 0x1f && bytes[1] == 0x8b {
        let mut decoder = flate2::read::GzDecoder::new(bytes);
        let mut out = Vec::new();
        decoder.read_to_end(&mut out).map_err(|e| anyhow!("gzip 解压失败: {}", e))?;
        Ok(out)
    } else {
        Ok(bytes.to_vec())
    }
}

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

    #[test]
    fn test_validate_cone_params() {
        // 合法范围：0~30°
        assert!(validate_cone_params(10.0, 41.0, 0.1).is_ok());
        assert!(validate_cone_params(10.0, 41.0, 5.0).is_ok()); // 建议上限内
        assert!(validate_cone_params(10.0, 41.0, 29.0).is_ok()); // 超建议但仍合法
        assert!(validate_cone_params(10.0, 41.0, 31.0).is_err()); // 超硬性上限
        assert!(validate_cone_params(10.0, 91.0, 0.1).is_err()); // dec 越界
    }

    #[test]
    fn test_decompress_if_gzip_plain() {
        let plain = b"SIMPLE  =                    T";
        let out = decompress_if_gzip(plain).unwrap();
        assert_eq!(out, plain);
    }

    #[test]
    fn test_decompress_if_gzip_compressed() {
        use flate2::write::GzEncoder;
        use flate2::Compression;
        use std::io::Write;
        let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
        encoder.write_all(b"hello fits").unwrap();
        let gz = encoder.finish().unwrap();
        let out = decompress_if_gzip(&gz).unwrap();
        assert_eq!(out, b"hello fits");
    }

    #[test]
    fn test_row_to_candidate() {
        let row = LamostSpectrumRow {
            obsid: 438809089, designation: None, ra_obs: Some(10.5), dec_obs: Some(41.0),
            z: None, class: None, subclass: None,
            sn_u: None, sn_g: None, sn_r: None, sn_i: None, sn_z: None,
        };
        let c = row_to_candidate(&row, "lrs");
        assert_eq!(c.source_id, "438809089");
        assert_eq!(c.ra.unwrap(), 10.5);
    }
}