AstroResearch/src/services/cds/vizier.rs

// src/services/cds/vizier.rs
//
// VizieR TAP 业务服务层 —— 合并缓存、几何查询、星表发现、导出为一个模块
//
// ┌──────────────────────────────────────────────────────────────┐
// │ 1. 缓存层        query_adql_cached / query_table            │
// │ 2. ADQL 安全     sanitize_identifier / escape_string_literal │
// │ 3. 几何查询      cone_search（nearest/all）                      │
// │ 4. 星表发现      VizierCatalog { search, describe, lookup }  │
// │ 5. 数据导出      rows_to_csv                                 │
// └──────────────────────────────────────────────────────────────┘
//
// 纯通信层（HTTP、VOTable/JSON 解析）在 clients::cds::vizier。
// 外部（api/、agent tools/、cli）统一通过本模块调用。

use crate::clients::ads::AdsClient;
use crate::clients::cds::vizier::{FieldInfo, VizierClient, VizierQueryResult};
use anyhow::{anyhow, Result};
use sha2::{Digest, Sha256};
use sqlx::SqlitePool;
use std::time::Duration;
use tracing::{error, info, warn};

// ═══════════════════════════════════════════════════════════════
// 1. 缓存层
// ═══════════════════════════════════════════════════════════════

const CACHE_TTL_SECS: i64 = 7 * 24 * 3600;

/// 执行 ADQL 查询（带缓存）
pub async fn query_adql_cached(
    pool: &SqlitePool,
    client: &VizierClient,
    adql: &str,
    max_records: i64,
) -> Result<VizierQueryResult> {
    let query_hash = hash_query(adql, max_records);

    if let Some(cached) = fetch_cache(pool, &query_hash).await? {
        info!("[VizieR] 缓存命中 (hash={:.12})", query_hash);
        return Ok(cached);
    }

    tokio::time::sleep(Duration::from_millis(50)).await;
    let result = client.run_adql(adql, max_records).await?;

    if let Err(e) = write_cache(pool, &query_hash, adql, max_records, &result).await {
        error!("[VizieR] 写入查询缓存失败: {}", e);
    }

    Ok(result)
}

/// 执行 ADQL 查询（不经过缓存）
pub async fn query_adql(
    client: &VizierClient,
    adql: &str,
    max_records: i64,
) -> Result<VizierQueryResult> {
    client.run_adql(adql, max_records).await
}

/// 便捷：按表名取行
pub async fn query_table(
    pool: &SqlitePool,
    client: &VizierClient,
    table_name: &str,
    columns: &[String],
    limit: i64,
) -> Result<VizierQueryResult> {
    let table = sanitize_identifier(table_name)?;
    let cols = if columns.is_empty() {
        "*".to_string()
    } else {
        columns
            .iter()
            .map(|c| sanitize_identifier(c))
            .collect::<Result<Vec<_>>>()?
            .join(", ")
    };
    let table_ref = if table.contains('/') || table.contains(' ') {
        format!("\"{}\"", table)
    } else {
        table
    };
    let adql = format!("SELECT TOP {} {} FROM {}", limit.max(1), cols, table_ref);
    query_adql_cached(pool, client, &adql, limit.max(1)).await
}

// ── 缓存辅助 ──

fn hash_query(adql: &str, max_records: i64) -> String {
    let mut hasher = Sha256::new();
    hasher.update(adql.as_bytes());
    hasher.update(max_records.to_le_bytes());
    format!("{:x}", hasher.finalize())
}

#[derive(sqlx::FromRow)]
struct CacheRow {
    result_json: String,
    expires_at: Option<chrono::DateTime<chrono::Utc>>,
}

async fn fetch_cache(pool: &SqlitePool, query_hash: &str) -> Result<Option<VizierQueryResult>> {
    let row = sqlx::query_as::<_, CacheRow>(
        "SELECT result_json, expires_at FROM vizier_query_cache WHERE query_hash = ? LIMIT 1",
    )
    .bind(query_hash)
    .fetch_optional(pool)
    .await
    .map_err(|e| anyhow!("查询 VizieR 缓存失败: {}", e))?;

    match row {
        Some(r) => {
            if let Some(exp) = r.expires_at {
                if exp < chrono::Utc::now() {
                    warn!("[VizieR] 缓存已过期 (hash={:.12})", query_hash);
                    return Ok(None);
                }
            }
            let result: VizierQueryResult = serde_json::from_str(&r.result_json)
                .map_err(|e| anyhow!("反序列化缓存失败: {}", e))?;
            Ok(Some(result))
        }
        None => Ok(None),
    }
}

async fn write_cache(
    pool: &SqlitePool,
    query_hash: &str,
    adql: &str,
    max_records: i64,
    result: &VizierQueryResult,
) -> Result<()> {
    let result_json = serde_json::to_string(result)?;
    let expires_at = chrono::Utc::now() + chrono::Duration::seconds(CACHE_TTL_SECS);
    sqlx::query(
        "INSERT OR REPLACE INTO vizier_query_cache (query_hash, adql, max_records, result_json, created_at, expires_at) VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP, ?)",
    )
    .bind(query_hash)
    .bind(adql)
    .bind(max_records)
    .bind(&result_json)
    .bind(expires_at)
    .execute(pool)
    .await
    .map_err(|e| anyhow!("写入 VizieR 缓存失败: {}", e))?;
    Ok(())
}

// ═══════════════════════════════════════════════════════════════
// 2. ADQL 安全
// ═══════════════════════════════════════════════════════════════

/// 校验 ADQL 标识符（表名/列名），白名单：字母、数字、_、.、/
pub fn sanitize_identifier(name: &str) -> Result<String> {
    let trimmed = name.trim();
    if trimmed.is_empty() {
        return Err(anyhow!("标识符不能为空"));
    }
    let valid = trimmed
        .chars()
        .all(|c| c.is_alphanumeric() || c == '_' || c == '.' || c == '/');
    if !valid {
        return Err(anyhow!(
            "标识符包含非法字符（仅允许字母、数字、下划线、点、斜杠）: {}",
            trimmed
        ));
    }
    Ok(trimmed.to_string())
}

/// 转义字符串字面量（单引号 → 双单引号）
pub fn escape_string_literal(s: &str) -> String {
    s.replace('\'', "''")
}

// ═══════════════════════════════════════════════════════════════
// 3. 几何查询（Cone Search / Cross-Match）
// ═══════════════════════════════════════════════════════════════

/// 锥形检索
///
/// strategy="nearest" 时按角距离排序，返回最近的源；strategy="all" 时无序返回全部。
pub async fn cone_search(
    pool: &SqlitePool,
    client: &VizierClient,
    ra: f64,
    dec: f64,
    radius_deg: f64,
    table_name: &str,
    max_records: i64,
    nearest: bool,
) -> Result<VizierQueryResult> {
    let table = sanitize_identifier(table_name)?;
    let table_ref = format!("\"{}\"", table);

    if !(0.0..=20.0).contains(&radius_deg) {
        return Err(anyhow!("检索半径应在 0~20 度之间，当前: {}", radius_deg));
    }
    if !(-360.0..=360.0).contains(&ra) || !(-90.0..=90.0).contains(&dec) {
        return Err(anyhow!("坐标范围异常 (ra={}, dec={})", ra, dec));
    }

    let adql = if nearest {
        format!(
            "SELECT TOP {} *, DISTANCE(POINT('ICRS', ra, dec), POINT('ICRS', {}, {})) AS separation_deg FROM {} WHERE 1=CONTAINS(POINT('ICRS', ra, dec), CIRCLE('ICRS', {}, {}, {})) ORDER BY separation_deg ASC",
            max_records.max(1), ra, dec, table_ref, ra, dec, radius_deg
        )
    } else {
        format!(
            "SELECT TOP {} * FROM {} WHERE 1=CONTAINS(POINT('ICRS', ra, dec), CIRCLE('ICRS', {}, {}, {}))",
            max_records.max(1), table_ref, ra, dec, radius_deg
        )
    };

    query_adql_cached(pool, client, &adql, max_records.max(1)).await
}

// ═══════════════════════════════════════════════════════════════
// 4. 星表发现（VizierCatalog 统一入口）
// ═══════════════════════════════════════════════════════════════

#[derive(Debug, Clone)]
pub struct CatalogEntry {
    pub table_name: String,
    pub description: String,
    pub schema_name: String,
    pub nrows: Option<i64>,
}

#[derive(Debug, Clone, serde::Serialize)]
pub struct ColumnInfo {
    pub column_name: String,
    pub datatype: String,
    pub unit: Option<String>,
    pub description: Option<String>,
    pub ucd: Option<String>,
}

#[derive(Debug, Clone)]
pub struct ExportResult {
    pub csv: String,
    pub row_count: usize,
    pub column_count: usize,
}

#[derive(Debug)]
pub struct ExportToFileResult {
    pub path: std::path::PathBuf,
    pub row_count: usize,
    pub column_count: usize,
    pub size_bytes: usize,
}

pub struct VizierCatalog<'a> {
    pub pool: &'a SqlitePool,
    pub vizier: &'a VizierClient,
    pub ads: Option<&'a AdsClient>,
}

impl<'a> VizierCatalog<'a> {
    pub fn new(pool: &'a SqlitePool, vizier: &'a VizierClient) -> Self {
        Self {
            pool,
            vizier,
            ads: None,
        }
    }

    pub fn with_ads(pool: &'a SqlitePool, vizier: &'a VizierClient, ads: &'a AdsClient) -> Self {
        Self {
            pool,
            vizier,
            ads: Some(ads),
        }
    }

    /// 按关键词搜索星表目录
    pub async fn search(&self, keyword: &str, limit: usize) -> Result<Vec<CatalogEntry>> {
        let kw_escaped = keyword.replace('\'', "''");
        let top = (limit * 2).min(100);
        let adql = format!(
            "SELECT TOP {top} table_name, description, schema_name, nrows \
             FROM TAP_SCHEMA.tables \
             WHERE table_name LIKE '%{kw}%' OR description LIKE '%{kw}%' \
             ORDER BY nrows DESC",
            top = top,
            kw = kw_escaped,
        );
        let result = query_adql_cached(self.pool, self.vizier, &adql, top as i64).await?;
        Ok(parse_catalog_rows(&result.rows, limit))
    }

    /// 查看表的列结构
    pub async fn describe(&self, table_name: &str) -> Result<Vec<ColumnInfo>> {
        let tn = table_name.replace('\'', "''");
        let adql = format!(
            "SELECT TOP 500 column_name, datatype, unit, description, ucd \
             FROM TAP_SCHEMA.columns WHERE table_name = '\"{tn}\"' OR table_name = '{tn}'",
        );
        let result = query_adql_cached(self.pool, self.vizier, &adql, 500).await?;
        Ok(parse_column_rows(&result.rows))
    }

    /// 下载数据并导出为 CSV
    pub async fn export(&self, adql: &str, max_records: i64) -> Result<ExportResult> {
        let result = query_adql_cached(self.pool, self.vizier, adql, max_records).await?;
        let csv = rows_to_csv(&result.fields, &result.rows);
        Ok(ExportResult {
            row_count: result.row_count,
            column_count: result.fields.len(),
            csv,
        })
    }

    /// 构建导出 ADQL（table + columns → ADQL）
    pub fn build_export_adql(table: &str, columns: &str, limit: i64) -> Result<String> {
        let table_ref = if table.contains('/') || table.contains(' ') {
            format!("\"{}\"", sanitize_identifier(table)?)
        } else {
            sanitize_identifier(table)?
        };
        Ok(format!("SELECT TOP {} {} FROM {}", limit, columns, table_ref))
    }

    /// 导出星表数据到 CSV 文件
    ///
    /// 优先使用 adql 参数；若未提供则从 table + columns 构建。
    /// 文件保存到 `{library_dir}/vizier/{catalog_name}/{timestamp}.csv`。
    pub async fn export_to_file(
        &self,
        library_dir: &str,
        adql: Option<&str>,
        table: Option<&str>,
        columns: Option<&str>,
        limit: i64,
    ) -> Result<ExportToFileResult> {
        let adql_str = if let Some(a) = adql {
            a.to_string()
        } else if let Some(t) = table {
            let cols = columns.unwrap_or("*");
            Self::build_export_adql(t, cols, limit)?
        } else {
            return Err(anyhow!("export 需要 'adql' 或 'table' 参数"));
        };

        info!(
            "[VizieR:export] ADQL: {}",
            adql_str.chars().take(150).collect::<String>()
        );

        let export = self.export(&adql_str, limit).await?;

        let catalog_name = table.unwrap_or("adql_query");
        let safe_name: String = catalog_name
            .chars()
            .map(|c| if c.is_alphanumeric() || c == '-' || c == '_' { c } else { '_' })
            .collect();

        let ts = chrono::Utc::now().format("%Y%m%d_%H%M%S");
        let output_path = std::path::PathBuf::from(library_dir)
            .join("vizier")
            .join(&safe_name)
            .join(format!("{}.csv", ts));

        if let Some(parent) = output_path.parent() {
            tokio::fs::create_dir_all(parent).await
                .map_err(|e| anyhow!("创建目录失败: {}", e))?;
        }
        tokio::fs::write(&output_path, &export.csv).await
            .map_err(|e| anyhow!("写入文件失败: {}", e))?;

        Ok(ExportToFileResult {
            path: output_path,
            row_count: export.row_count,
            column_count: export.column_count,
            size_bytes: export.csv.len(),
        })
    }

    /// 通过文献 bibcode 查找关联的 VizieR 表
    pub async fn lookup(&self, bibcode: &str, limit: usize) -> Result<Vec<CatalogEntry>> {
        let ads = self.ads.ok_or_else(|| anyhow!("lookup 需要 ADS 客户端"))?;

        let query = format!("bibcode:{}", bibcode);
        let docs = ads.search(&query, 0, 1, "score desc").await?;
        let doc = docs
            .first()
            .ok_or_else(|| anyhow!("ADS 未找到 bibcode: {}", bibcode))?;

        let data_links = match doc.data {
            Some(ref links) if !links.is_empty() => links,
            _ => return Ok(Vec::new()),
        };

        let mut entries = Vec::new();
        for link in data_links {
            if let Some(table_name) = extract_table_from_ads_url(link) {
                let adql = format!(
                    "SELECT TOP 1 table_name, description, schema_name, nrows \
                     FROM TAP_SCHEMA.tables WHERE table_name = '\"{}\"' OR table_name = '{}'",
                    table_name.replace('"', ""),
                    table_name.replace('\'', "''"),
                );
                if let Ok(result) = query_adql_cached(self.pool, self.vizier, &adql, 1).await {
                    entries.extend(parse_catalog_rows(&result.rows, 1));
                }
            }
        }

        entries.truncate(limit);
        Ok(entries)
    }
}

fn parse_catalog_rows(rows: &[Vec<serde_json::Value>], limit: usize) -> Vec<CatalogEntry> {
    rows.iter()
        .filter_map(|row| {
            if row.len() < 4 {
                return None;
            }
            Some(CatalogEntry {
                table_name: row[0].as_str()?.trim_matches('"').to_string(),
                description: row[1].as_str().unwrap_or("").to_string(),
                schema_name: row[2].as_str().unwrap_or("").to_string(),
                nrows: row[3].as_f64().map(|v| v as i64),
            })
        })
        .take(limit)
        .collect()
}

fn parse_column_rows(rows: &[Vec<serde_json::Value>]) -> Vec<ColumnInfo> {
    rows.iter()
        .filter_map(|row| {
            if row.len() < 5 {
                return None;
            }
            Some(ColumnInfo {
                column_name: row[0].as_str()?.trim_matches('"').to_string(),
                datatype: row[1].as_str().unwrap_or("unknown").to_string(),
                unit: row[2].as_str().map(|s| s.to_string()),
                description: row[3].as_str().map(|s| s.to_string()),
                ucd: row[4].as_str().map(|s| s.to_string()),
            })
        })
        .collect()
}

// ═══════════════════════════════════════════════════════════════
// 5. 数据导出
// ═══════════════════════════════════════════════════════════════

pub fn rows_to_csv(columns: &[FieldInfo], rows: &[Vec<serde_json::Value>]) -> String {
    let mut csv = String::new();
    let headers: Vec<&str> = columns.iter().map(|f| f.name.as_str()).collect();
    csv.push_str(&headers.join(","));
    csv.push('\n');
    for row in rows {
        let cells: Vec<String> = row
            .iter()
            .map(|v| match v {
                serde_json::Value::Null => String::new(),
                serde_json::Value::String(s) => {
                    if s.contains(',') || s.contains('\n') || s.contains('"') {
                        format!("\"{}\"", s.replace('"', "\"\""))
                    } else {
                        s.clone()
                    }
                }
                other => other.to_string(),
            })
            .collect();
        csv.push_str(&cells.join(","));
        csv.push('\n');
    }
    csv
}

/// 从 ADS data 链接中提取 VizieR 表名
pub fn extract_table_from_ads_url(url: &str) -> Option<String> {
    if let Some(pos) = url.find("-source=") {
        let table = &url[pos + 8..];
        let table = table.split_whitespace().next()?;
        let table = table.trim_matches('"').trim_matches('\'');
        return Some(table.to_string());
    }
    if let Some(pos) = url.find("VizieR") {
        let after = &url[pos..];
        if let Some(qpos) = after.find('?') {
            let query = &after[qpos + 1..];
            for param in query.split('&') {
                if let Some(val) = param.strip_prefix("-source=") {
                    return Some(val.to_string());
                }
            }
        }
    }
    if url.contains('/') && !url.starts_with("http") {
        return Some(url.to_string());
    }
    None
}

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

    #[test]
    fn test_sanitize_identifier_valid() {
        assert_eq!(sanitize_identifier("gaiadr3").unwrap(), "gaiadr3");
        assert_eq!(
            sanitize_identifier("I/355/gaiadr3").unwrap(),
            "I/355/gaiadr3"
        );
    }

    #[test]
    fn test_sanitize_identifier_rejects() {
        assert!(sanitize_identifier("ra; DROP TABLE").is_err());
        assert!(sanitize_identifier("' OR '1'='1").is_err());
    }

    #[test]
    fn test_escape_string_literal() {
        assert_eq!(escape_string_literal("O'Brien"), "O''Brien");
    }

    #[test]
    fn test_hash_query_stable() {
        let h1 = hash_query("SELECT 1", 10);
        let h2 = hash_query("SELECT 1", 10);
        let h3 = hash_query("SELECT 1", 20);
        assert_eq!(h1, h2);
        assert_ne!(h1, h3);
    }

    #[test]
    fn test_extract_table_from_ads_url() {
        assert_eq!(
            extract_table_from_ads_url(
                "https://vizier.cds.unistra.fr/viz-bin/VizieR?-source=I/355/gaiadr3"
            ),
            Some("I/355/gaiadr3".to_string())
        );
        assert_eq!(
            extract_table_from_ads_url("I/355/gaiadr3"),
            Some("I/355/gaiadr3".to_string())
        );
        assert_eq!(extract_table_from_ads_url("http://example.com"), None);
    }

    #[tokio::test]
    #[ignore = "需要网络访问"]
    async fn test_live_vizier_tap() {
        let result = query_adql(
            &VizierClient::new("https://tapvizier.cds.unistra.fr/TAPVizieR/tap", 60).unwrap(),
            r#"SELECT TOP 3 DR3Name, RA_ICRS FROM "I/355/gaiadr3""#,
            3,
        )
        .await
        .unwrap();
        assert_eq!(result.row_count, 3);
    }
}