//! 数据导出：Axum 流式路由（大文件不走 JSON 序列化）。
//!
//! 鉴权镜像 [`crate::api::upload`]：从 cookie 取 session，校验管理员。
//! 两种来源：按表导出（表名白名单）/ 按查询导出（强制只读，AST 校验）。
//! 两种格式：CSV（COPY TO STDOUT 流式）/ SQL（逐行拼 INSERT）。

#![allow(clippy::unused_unit)]

use axum::body::Body;
use axum::extract::Query;
use axum::http::{header, HeaderMap, StatusCode};
use axum::response::IntoResponse;
use futures::StreamExt;
use serde::Deserialize;

use crate::auth::session::parse_session_token;

#[derive(Deserialize)]
pub struct ExportParams {
    /// `table:<name>` 或 `query:<SELECT ...>`。
    pub source: String,
    /// `sql` 或 `csv`。
    pub format: String,
    /// 是否带表头（CSV）/ 列名（SQL INSERT）。默认 true。
    pub include_columns: Option<bool>,
}

/// POST/GET /api/database/export —— 流式导出。
pub async fn export_data(
    headers: HeaderMap,
    Query(params): Query<ExportParams>,
) -> Result<impl IntoResponse, (StatusCode, String)> {
    // 1. 鉴权：cookie → session → admin
    let cookie_header = headers
        .get("cookie")
        .and_then(|h| h.to_str().ok())
        .unwrap_or("");
    let token = parse_session_token(cookie_header).map(str::to_string);
    let token = match token {
        Some(t) => t,
        None => return Err((StatusCode::UNAUTHORIZED, "未登录".to_string())),
    };
    let user = match crate::api::auth::get_user_by_token(&token).await {
        Ok(Some(u)) => u,
        _ => return Err((StatusCode::UNAUTHORIZED, "会话已过期".to_string())),
    };
    if user.role != crate::models::user::UserRole::Admin {
        return Err((StatusCode::FORBIDDEN, "权限不足".to_string()));
    }

    // 2. 解析来源 + 白名单/只读校验
    let include_columns = params.include_columns.unwrap_or(true);
    let (source_sql, table_name) = parse_source(&params.source)?;

    match params.format.as_str() {
        "csv" => export_csv(source_sql, include_columns).await,
        "sql" => export_sql(source_sql, include_columns, table_name).await,
        _ => Err((StatusCode::BAD_REQUEST, "不支持的格式".to_string())),
    }
    .map(|(body, content_type, filename)| {
        let disposition = format!("attachment; filename=\"{}\"", filename);
        let disposition_value = axum::http::HeaderValue::from_str(&disposition)
            .unwrap_or_else(|_| axum::http::HeaderValue::from_static("attachment"));
        let content_type_value = axum::http::HeaderValue::from_str(content_type)
            .unwrap_or_else(|_| axum::http::HeaderValue::from_static("application/octet-stream"));
        (
            StatusCode::OK,
            [
                (header::CONTENT_TYPE, content_type_value),
                (header::CONTENT_DISPOSITION, disposition_value),
            ],
            body,
        )
    })
}

/// 解析导出来源，返回（内部 SQL，表名）。
/// - `table:posts` → 校验表名合法后，`SELECT * FROM "posts"`（只读）+ 表名 "posts"。
/// - `query:SELECT ...` → 校验为只读语句后原样返回，表名为 "export"。
fn parse_source(source: &str) -> Result<(String, String), (StatusCode, String)> {
    if let Some(table) = source.strip_prefix("table:") {
        // 表名白名单：仅允许标识符字符，防注入
        let t = table.trim();
        if t.is_empty() || !is_simple_ident(t) {
            return Err((StatusCode::BAD_REQUEST, "无效的表名".to_string()));
        }
        Ok((format!("SELECT * FROM \"{}\"", t), t.to_string()))
    } else if let Some(query) = source.strip_prefix("query:") {
        // 只读校验：sqlparser 解析后所有语句均为 Query/Explain
        let dialect = sqlparser::dialect::PostgreSqlDialect {};
        let parsed = sqlparser::parser::Parser::parse_sql(&dialect, query)
            .map_err(|_| (StatusCode::BAD_REQUEST, "SQL 解析失败".to_string()))?;
        if parsed.iter().any(|s| !is_read_only_ast(s)) {
            return Err((
                StatusCode::BAD_REQUEST,
                "导出查询必须是只读（SELECT/EXPLAIN）".to_string(),
            ));
        }
        Ok((query.to_string(), "export".to_string()))
    } else {
        Err((
            StatusCode::BAD_REQUEST,
            "source 必须是 table:<name> 或 query:<sql>".to_string(),
        ))
    }
}

/// 简单标识符校验（字母数字下划线，防 SQL 注入与路径穿越）。
fn is_simple_ident(s: &str) -> bool {
    !s.is_empty() && s.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
}

/// 判断 AST 是否只读（SELECT/EXPLAIN）。
fn is_read_only_ast(stmt: &sqlparser::ast::Statement) -> bool {
    use sqlparser::ast::Statement;
    matches!(stmt, Statement::Query(_) | Statement::Explain { .. })
}

/// CSV 导出：用 COPY ... TO STDOUT WITH CSV 流式（大表不 OOM）。
async fn export_csv(
    source_sql: String,
    include_columns: bool,
) -> Result<(Body, &'static str, String), (StatusCode, String)> {
    let client = crate::db::pool::get_conn()
        .await
        .map_err(|_| (StatusCode::SERVICE_UNAVAILABLE, "数据库不可用".to_string()))?;

    let header_clause = if include_columns { "HEADER" } else { "" };
    let copy_stmt = format!("COPY ({}) TO STDOUT WITH CSV {}", source_sql, header_clause);
    let stream = client
        .copy_out(&copy_stmt)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("COPY 失败：{e}")))?;

    // tokio_postgres 的 copy_out 流产出 Bytes，直接转 axum Body。
    let mapped = stream.map(|res| res.map_err(std::io::Error::other));
    let body = Body::from_stream(mapped);
    Ok((body, "text/csv; charset=utf-8", "export.csv".to_string()))
}

/// SQL 导出：逐行拼 INSERT（纯数据，不含 DDL）。
async fn export_sql(
    source_sql: String,
    include_columns: bool,
    table_name: String,
) -> Result<(Body, &'static str, String), (StatusCode, String)> {
    let client = crate::db::pool::get_conn()
        .await
        .map_err(|_| (StatusCode::SERVICE_UNAVAILABLE, "数据库不可用".to_string()))?;

    let rows = client
        .query(&source_sql, &[])
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("查询失败：{e}")))?;

    let columns: Vec<String> = rows
        .first()
        .map(|r| r.columns().iter().map(|c| c.name().to_string()).collect())
        .unwrap_or_default();

    let mut out = String::new();
    out.push_str("-- Yggdrasil 数据导出（仅数据，不含 schema）\n");
    let col_clause = if include_columns && !columns.is_empty() {
        format!("({})", columns.join(", "))
    } else {
        String::new()
    };

    for r in &rows {
        let vals: Vec<String> = (0..r.len()).map(|i| sql_quote_cell(r, i)).collect();
        out.push_str(&format!(
            "INSERT INTO {} {} VALUES ({});\n",
            table_name,
            col_clause,
            vals.join(", ")
        ));
    }

    let body = Body::from(out);
    Ok((
        body,
        "application/sql; charset=utf-8",
        "export.sql".to_string(),
    ))
}

/// 把一个单元格值转成 SQL 字面量（字符串单引号转义，数字原样，NULL）。
fn sql_quote_cell(row: &tokio_postgres::Row, idx: usize) -> String {
    let ty = row
        .columns()
        .get(idx)
        .map(|c| c.type_().name())
        .unwrap_or("");
    match ty {
        "int2" => row
            .try_get::<_, Option<i16>>(idx)
            .ok()
            .flatten()
            .map(|v| v.to_string())
            .unwrap_or_else(|| "NULL".into()),
        "int4" => row
            .try_get::<_, Option<i32>>(idx)
            .ok()
            .flatten()
            .map(|v| v.to_string())
            .unwrap_or_else(|| "NULL".into()),
        "int8" => row
            .try_get::<_, Option<i64>>(idx)
            .ok()
            .flatten()
            .map(|v| v.to_string())
            .unwrap_or_else(|| "NULL".into()),
        "float4" | "float8" => row
            .try_get::<_, Option<f64>>(idx)
            .ok()
            .flatten()
            .map(|v| v.to_string())
            .unwrap_or_else(|| "NULL".into()),
        "bool" => row
            .try_get::<_, Option<bool>>(idx)
            .ok()
            .flatten()
            .map(|v| if v { "TRUE" } else { "FALSE" }.to_string())
            .unwrap_or_else(|| "NULL".into()),
        _ => {
            // 字符串/时间戳等：单引号转义
            match row.try_get::<_, Option<String>>(idx) {
                Ok(Some(s)) => format!("'{}'", s.replace('\'', "''")),
                _ => "NULL".into(),
            }
        }
    }
}

#[cfg(all(test, feature = "server"))]
mod tests {
    use super::*;

    // ---- is_simple_ident：表名白名单（防 SQL 注入与路径穿越）----

    #[test]
    fn simple_ident_accepts_legal_names() {
        assert!(is_simple_ident("posts"));
        assert!(is_simple_ident("user_posts"));
        assert!(is_simple_ident("table1"));
        assert!(is_simple_ident("a"));
        assert!(is_simple_ident("_private"));
        assert!(is_simple_ident("ABC123xyz"));
    }

    #[test]
    fn simple_ident_rejects_empty_and_whitespace() {
        assert!(!is_simple_ident(""));
        assert!(!is_simple_ident("   "));
    }

    #[test]
    fn simple_ident_rejects_injection_vectors() {
        // SQL 注入：引号、分号、注释、空格分隔
        assert!(!is_simple_ident("posts; DROP TABLE users"));
        assert!(!is_simple_ident("posts' OR '1'='1"));
        assert!(!is_simple_ident("posts--"));
        assert!(!is_simple_ident("public.posts"));
        assert!(!is_simple_ident("posts where 1=1"));
        // 路径穿越
        assert!(!is_simple_ident("../etc/passwd"));
        assert!(!is_simple_ident("..\\windows"));
        assert!(!is_simple_ident("/etc/passwd"));
        // 反引号 / 双引号标识符语法
        assert!(!is_simple_ident("`posts`"));
        assert!(!is_simple_ident("\"posts\""));
        // 连字符（常见合法表名片段，但本白名单禁止，防 `a-b` 被 SQL 解析为减法）
        assert!(!is_simple_ident("my-posts"));
    }

    #[test]
    fn simple_ident_rejects_unicode_and_non_ascii() {
        assert!(!is_simple_ident("文章"));
        assert!(!is_simple_ident("posts²"));
        assert!(!is_simple_ident("café"));
    }

    // ---- parse_source：来源解析 + 只读校验（安全入口）----

    #[test]
    fn parse_table_source_builds_select_star() {
        let (sql, name) = parse_source("table:posts").unwrap();
        assert_eq!(sql, "SELECT * FROM \"posts\"");
        assert_eq!(name, "posts");
    }

    #[test]
    fn parse_table_source_trims_whitespace() {
        let (sql, name) = parse_source("table:  posts  ").unwrap();
        assert_eq!(sql, "SELECT * FROM \"posts\"");
        assert_eq!(name, "posts");
    }

    #[test]
    fn parse_table_source_rejects_bad_name() {
        // 非法表名（注入/穿越/空）必须在拼进 SQL 前被拒，防 SQL 注入。
        for bad in [
            "table:",
            "table:   ",
            "table:posts; DROP TABLE users",
            "table:../secret",
            "table:a b",
            "table:\"x\"",
        ] {
            let (code, msg) = parse_source(bad).expect_err(bad);
            assert_eq!(code, StatusCode::BAD_REQUEST);
            assert!(msg.contains("表名"), "{bad}: {msg}");
        }
    }

    #[test]
    fn parse_query_source_accepts_select_and_explain() {
        assert!(parse_source("query:SELECT * FROM posts").is_ok());
        assert!(parse_source("query:SELECT id, title FROM posts WHERE id > 0").is_ok());
        assert!(parse_source("query:EXPLAIN SELECT * FROM posts").is_ok());
        assert!(parse_source("query:SELECT 1").is_ok());
        // 复杂但只读的查询
        assert!(parse_source("query:WITH t AS (SELECT 1) SELECT * FROM t").is_ok());
    }

    #[test]
    fn parse_query_source_rejects_all_write_operations() {
        // 每一种写/破坏操作都必须被 AST 校验拦截——这是数据导出的核心安全不变量。
        for evil in [
            "query:INSERT INTO posts VALUES (1)",
            "query:UPDATE posts SET title='x'",
            "query:DELETE FROM posts",
            "query:DROP TABLE posts",
            "query:TRUNCATE posts",
            "query:CREATE TABLE evil (id int)",
            "query:ALTER TABLE posts ADD COLUMN x int",
            "query:GRANT SELECT ON posts TO public",
            // 即便嵌在 SELECT 里，子查询的写操作也要拒
            "query:INSERT INTO logs VALUES (1) RETURNING 1",
        ] {
            let (code, msg) = parse_source(evil).expect_err(evil);
            assert_eq!(code, StatusCode::BAD_REQUEST, "{evil}");
            // 写操作要么被 AST 校验拒（"只读"），要么 sqlparser 解析阶段就失败
            assert!(
                msg.contains("只读") || msg.contains("解析失败"),
                "{evil}: {msg}"
            );
        }
    }

    #[test]
    fn parse_query_source_rejects_unparseable_sql() {
        for bad in [
            "query:这不是SQL",
            "query:SELECT FROM",
            "query:@#$%",
            "query:1+1",
        ] {
            assert!(parse_source(bad).is_err(), "{bad} 应解析失败或被拒");
        }
    }

    #[test]
    fn parse_query_source_allows_empty_result_set() {
        // sqlparser 容忍纯分号（解析为零语句），`any()` 在空集上为 false → 放行。
        // 锁定该契约：这只产生空导出，无数据泄露，属可接受行为；
        // 若未来想收紧为"空查询拒绝"，此测试会提醒你更新。
        assert!(parse_source("query:;;;").is_ok());
    }

    #[test]
    fn parse_source_rejects_unknown_prefix() {
        let (code, msg) = parse_source("unknown:posts").expect_err("未知前缀");
        assert_eq!(code, StatusCode::BAD_REQUEST);
        assert!(msg.contains("source"));
        // 完全无前缀的裸字符串也应被拒
        assert!(parse_source("posts").is_err());
        assert!(parse_source("").is_err());
    }
}
