from datetime import datetime
from typing import Literal

from fastapi import APIRouter, Query
from fastapi.responses import StreamingResponse
from io import BytesIO

from app.services.export_service import generate_excel, generate_csv

router = APIRouter()

# 存储最近的查询结果 (简化实现, 生产环境应使用 Redis 等缓存)
_cached_data: list = []


def set_export_data(data: list):
    """设置导出数据缓存."""
    global _cached_data
    _cached_data = data


def get_export_data() -> list:
    """获取导出数据缓存."""
    return _cached_data


@router.get("/export")
async def export_data(
    format: Literal["xlsx", "csv"] = Query("xlsx", description="导出格式"),
):
    """
    导出查询结果.

    Args:
        format: 导出格式 (xlsx 或 csv)

    Returns:
        文件下载响应
    """
    data = get_export_data()

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    if format == "xlsx":
        content = generate_excel(data)
        media_type = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
        filename = f"kol_data_{timestamp}.xlsx"
    else:
        content = generate_csv(data)
        media_type = "text/csv; charset=utf-8"
        filename = f"kol_data_{timestamp}.csv"

    return StreamingResponse(
        BytesIO(content),
        media_type=media_type,
        headers={
            "Content-Disposition": f'attachment; filename="{filename}"',
        },
    )