AcidMap/app/api/water.py

from fastapi import APIRouter, Form, HTTPException, BackgroundTasks, Query
from fastapi.responses import FileResponse, JSONResponse
import os
import logging
import tempfile
import shutil
from typing import Dict, Any, Optional
from datetime import datetime
import json
from pathlib import Path

# 导入服务层函数
from ..services.water_service import (
    process_land_to_visualization,
    get_land_statistics,
    get_base_dir
)

# 配置日志
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)

router = APIRouter()


# 清理临时路径的函数
def cleanup_temp_path(path: str) -> None:
    """删除临时文件或整个临时目录"""
    try:
        if os.path.exists(path):
            if os.path.isfile(path):
                os.unlink(path)
            elif os.path.isdir(path):
                shutil.rmtree(path)
    except Exception as e:
        logger.warning(f"清理临时路径失败: {path}, 错误: {e}")


@router.get("/default-map",
            summary="获取默认地图图片",
            description="返回指定土地类型的默认地图图片")
async def get_default_map(
        land_type: str = Query("水田", description="土地类型，如：'水田'、'旱地'或'水浇地'")
) -> FileResponse:
    """返回默认地图图片"""
    try:
        logger.info(f"获取默认地图: {land_type}")

        # 获取基础目录
        base_dir = get_base_dir()
        raster_dir = os.path.join(base_dir, "..", "static", "water", "Raster")
        map_path = os.path.join(raster_dir, f"{land_type}_Cd含量地图.jpg")

        if not os.path.exists(map_path):
            logger.warning(f"默认地图文件不存在: {map_path}")
            raise HTTPException(status_code=404, detail="地图文件不存在")

        return FileResponse(
            path=map_path,
            filename=f"{land_type}_Cd含量地图.jpg",
            media_type="image/jpeg"
        )
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取默认地图失败: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"获取默认地图失败: {str(e)}"
        )


@router.get("/default-histogram",
            summary="获取默认直方图",
            description="返回指定土地类型的默认直方图图片")
async def get_default_histogram(
        land_type: str = Query("水田", description="土地类型，如：'水田'、'旱地'或'水浇地'")
) -> FileResponse:
    """返回默认直方图图片"""
    try:
        logger.info(f"获取默认直方图: {land_type}")

        # 获取基础目录
        base_dir = get_base_dir()
        raster_dir = os.path.join(base_dir, "..", "static", "water", "Raster")
        hist_path = os.path.join(raster_dir, f"{land_type}_Cd含量直方图.jpg")

        if not os.path.exists(hist_path):
            logger.warning(f"默认直方图文件不存在: {hist_path}")
            raise HTTPException(status_code=404, detail="直方图文件不存在")

        return FileResponse(
            path=hist_path,
            filename=f"{land_type}_Cd含量直方图.jpg",
            media_type="image/jpeg"
        )
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取默认直方图失败: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"获取默认直方图失败: {str(e)}"
        )


@router.post("/calculate",
             summary="重新计算土地数据",
             description="根据输入的土地类型和系数重新计算土地数据")
async def recalculate_land_data(
        background_tasks: BackgroundTasks,
        land_type: str = Form(..., description="土地类型，如：'水田'、'旱地'或'水浇地'"),
        param1: float = Form(711, description="土地类型系数的第一个参数"),
        param2: float = Form(0.524, description="土地类型系数的第二个参数"),
        color_map_name: str = Form("绿-黄-红-紫", description="使用的色彩方案"),
        output_size: int = Form(8, description="输出图片的尺寸")
) -> Dict[str, Any]:
    """重新计算土地数据并返回结果路径"""
    try:
        logger.info(f"重新计算土地数据: {land_type}")

        # 获取默认目录
        base_dir = get_base_dir()
        raster_dir = os.path.join(base_dir, "..", "static", "water", "Raster")

        # 确保目录存在
        os.makedirs(raster_dir, exist_ok=True)

        # 构建系数参数
        coefficient_params = {
            land_type: (param1, param2)
        }

        # 调用服务层函数处理土地数据
        results = process_land_to_visualization(
            land_type=land_type,
            coefficient_params=coefficient_params,
            color_map_name=color_map_name,
            output_size=output_size
        )

        if not results or any(r is None for r in results):
            logger.error(f"重新计算土地数据失败: {land_type}")
            raise HTTPException(
                status_code=500,
                detail=f"重新计算土地数据失败: {land_type}"
            )

        cleaned_csv, shapefile, tif_file, map_output, hist_output, used_coeff = results

        # 定义默认路径
        default_map_path = os.path.join(raster_dir, f"{land_type}_Cd含量地图.jpg")
        default_hist_path = os.path.join(raster_dir, f"{land_type}_Cd含量直方图.jpg")

        # 移动文件到默认目录（覆盖旧文件）
        shutil.move(map_output, default_map_path)
        shutil.move(hist_output, default_hist_path)

        return {
            "map_path": default_map_path,
            "histogram_path": default_hist_path,
            "used_coeff": used_coeff
        }
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"重新计算土地数据失败: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"重新计算土地数据失败: {str(e)}"
        )


@router.get("/statistics",
            summary="获取土地类型统计数据",
            description="返回指定土地类型的Cd预测结果统计信息")
async def get_land_statistics_endpoint(
        land_type: str = Query("水田", description="土地类型，如：'水田'、'旱地'或'水浇地'")
) -> JSONResponse:
    """返回土地类型Cd预测结果的统计信息"""
    try:
        logger.info(f"获取土地类型统计数据: {land_type}")

        # 调用服务层函数获取统计数据
        stats = get_land_statistics(land_type)

        if not stats:
            logger.warning(f"未找到{land_type}的土地类型统计数据")
            raise HTTPException(status_code=404, detail="统计数据不存在")

        return JSONResponse(content=stats)
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取土地类型统计数据失败: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"获取土地类型统计数据失败: {str(e)}"
        )