CapacityReport/app/api/routers/database.py

import re
from datetime import datetime
from typing import Optional

import pandas as pd
from fastapi import APIRouter, Body, HTTPException
from fastapi.responses import FileResponse
from starlette.background import BackgroundTask

from app import state
from app.config import CACHE_DIR
from app.database import DatabaseManager
from app.services.platform import make_client
from app.utils.files import remove_file_safely
from app.warehouse import make_warehouse


router = APIRouter(tags=["database"])
INVALID_SHEET_NAME_CHARS = re.compile(r"[:\\/?*\[\]]")


def _resolve_requested_tables(
    table_name: Optional[str],
    table_names: Optional[list[str]],
) -> list[str]:
    names = table_names if table_names is not None else ([table_name] if table_name else [])
    return [name.strip() for name in names if isinstance(name, str) and name.strip()]


def _make_sheet_name(table_name: str, used_names: set[str]) -> str:
    base = INVALID_SHEET_NAME_CHARS.sub("_", table_name).strip("'").strip() or "Sheet"
    base = base[:31]
    sheet_name = base
    index = 2

    while sheet_name in used_names:
        suffix = f"_{index}"
        sheet_name = f"{base[:31 - len(suffix)]}{suffix}" or f"Sheet_{index}"
        index += 1

    used_names.add(sheet_name)
    return sheet_name


def _dataframe_from_table(db: DatabaseManager, table_name: str) -> pd.DataFrame:
    result = db.query_table(table_name, page=1, page_size=1000000)
    table_info = db.get_table_info(table_name)
    columns = [str(column["Field"]) for column in table_info["columns"]]
    return pd.DataFrame(result["data"], columns=columns)


def _db():
    """Direct MySQL DatabaseManager, or a Metrix-backed warehouse with the same interface."""
    return make_warehouse(state.current_config())


@router.post("/api/database/test")
def test_database():
    db = _db()
    success, message = db.test_connection()
    return {"success": success, "message": message}


@router.get("/api/database/info")
def get_database_info():
    db = _db()
    try:
        return {"success": True, **db.get_server_info()}
    except Exception as exc:
        return {"success": False, "error": str(exc)}


@router.get("/api/database/tables")
@router.post("/api/database/tables")
def get_tables():
    db = _db()
    try:
        return {"tables": db.get_tables()}
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc


@router.post("/api/database/table/info")
def get_table_info(table_name: str = Body(..., embed=True)):
    db = _db()
    try:
        return db.get_table_info(table_name)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc


@router.post("/api/database/table/data")
def query_table_data(
    table_name: str = Body(..., embed=True),
    page: int = Body(1),
    page_size: int = Body(50),
    order_by: Optional[str] = Body(None),
    order_dir: str = Body("ASC"),
):
    db = _db()
    try:
        return db.query_table(table_name, page, page_size, order_by=order_by, order_dir=order_dir)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc


@router.post("/api/database/table/query")
def query_table_with_filter(
    table_name: str = Body(..., embed=True),
    page: int = Body(1),
    page_size: int = Body(50),
    filters: Optional[dict[str, str]] = Body(None),
    order_by: Optional[str] = Body(None),
    order_dir: str = Body("ASC"),
):
    db = _db()
    try:
        return db.query_table(
            table_name,
            page,
            page_size,
            filters=filters or {},
            order_by=order_by,
            order_dir=order_dir,
        )
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc


@router.post("/api/database/table/truncate")
def truncate_table(table_name: str = Body(..., embed=True)):
    db = _db()
    try:
        db.truncate_table(table_name)
        return {"success": True, "message": f"表 {table_name} 已清空"}
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc


@router.post("/api/database/table/drop")
def drop_table(table_name: str = Body(..., embed=True)):
    db = _db()
    try:
        db.drop_table(table_name)
        return {"success": True, "message": f"表 {table_name} 已删除"}
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc


@router.post("/api/database/table/drop-all")
def drop_all_tables():
    db = _db()
    try:
        result = db.drop_all_tables()
        return {
            "success": True,
            "message": f"已删除 {result['dropped_count']} 个表",
            "dropped_count": result["dropped_count"],
            "tables": result["tables"],
        }
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc


@router.post("/api/database/execute")
def execute_sql(sql: str = Body(..., embed=True)):
    db = _db()
    try:
        success, result = db.execute_sql(sql)
        if success:
            return {"success": True, "result": result}
        raise HTTPException(status_code=400, detail=result)
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc


# Sync def so FastAPI runs it in a threadpool: exporting large tables (and the Metrix
# export-job polling) is blocking and would otherwise freeze the single-worker event loop.
@router.post("/api/download")
def download_table(
    table_name: Optional[str] = Body(None, embed=True),
    table_names: Optional[list[str]] = Body(None, embed=True),
    file_format: str = Body("csv", alias="format"),
):
    if file_format not in {"csv", "xlsx"}:
        raise HTTPException(status_code=400, detail="不支持的导出格式")

    requested_tables = _resolve_requested_tables(table_name, table_names)
    if not requested_tables:
        raise HTTPException(status_code=400, detail="请选择要导出的数据表")
    if file_format == "csv" and len(requested_tables) != 1:
        raise HTTPException(status_code=400, detail="CSV 每次只能导出一张表")

    config = state.current_config()
    if config.warehouse_type == "metrix":
        return _download_via_metrix(config, requested_tables, file_format)

    db = _db()
    try:
        available_tables = set(db.get_tables())
        missing_tables = [name for name in requested_tables if name not in available_tables]
        if missing_tables:
            raise HTTPException(status_code=400, detail=f"数据表不存在: {', '.join(missing_tables)}")

        table_frames = {
            name: _dataframe_from_table(db, name)
            for name in requested_tables
        }
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename_prefix = requested_tables[0] if len(requested_tables) == 1 else "tables"
    filename = f"{filename_prefix}_{timestamp}.{file_format}"
    filepath = CACHE_DIR / filename
    CACHE_DIR.mkdir(parents=True, exist_ok=True)

    try:
        if file_format == "csv":
            table_frames[requested_tables[0]].to_csv(filepath, index=False, encoding="utf-8-sig")
            media_type = "text/csv"
        else:
            used_sheet_names: set[str] = set()
            with pd.ExcelWriter(filepath) as writer:
                for name, df in table_frames.items():
                    df.to_excel(writer, sheet_name=_make_sheet_name(name, used_sheet_names), index=False)
            media_type = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
    except Exception:
        remove_file_safely(filepath)
        raise

    return FileResponse(
        path=str(filepath),
        filename=filename,
        media_type=media_type,
        background=BackgroundTask(remove_file_safely, filepath),
    )


def _download_via_metrix(config, requested_tables: list[str], file_format: str) -> FileResponse:
    """Metrix 仓库模式：用平台导出任务生成文件后流式返回（避免分页上限丢行）。"""
    metrix = config.metrix.normalized()
    client = make_client(metrix)
    try:
        job_id = client.submit_export(metrix.database_conn_id, requested_tables, file_format, metrix.target_database)
        job = client.wait_job(job_id)
        if job.get("status") != "success":
            raise HTTPException(status_code=500, detail=f"导出失败: {job.get('error_code') or job.get('status')}")
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename_prefix = requested_tables[0] if len(requested_tables) == 1 else "tables"
        filename = f"{filename_prefix}_{timestamp}.{file_format}"
        filepath = CACHE_DIR / filename
        CACHE_DIR.mkdir(parents=True, exist_ok=True)
        client.download_job_file(job_id, filepath)
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc

    media_type = "text/csv" if file_format == "csv" else "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
    return FileResponse(
        path=str(filepath),
        filename=filename,
        media_type=media_type,
        background=BackgroundTask(remove_file_safely, filepath),
    )