geMoldInsight/src/api/routes.py

# api/routes.py
from fastapi import APIRouter, UploadFile, File, HTTPException, Request, Depends
from typing import Optional, Dict, Any, List
import uuid
from datetime import datetime
import os
from pathlib import Path

from models.schemas import ProcessingStatus, create_task_info
from utils.file_handler import FileHandler
from services.storage_integration_rustfs import StorageIntegrationService
from services.redis_task_manager import redis_task_manager
from services.task_query_service import TaskQueryService
from database.database import get_db_session

try:
    from celery_tasks import process_stp_task
    _use_celery = True
except ImportError:
    process_stp_task = None
    _use_celery = False
from utils.logger import get_logger
from sqlalchemy.ext.asyncio import AsyncSession
from core.cavity_layout_optimizer import CavityLayoutOptimizer
from core.mold_system_designer import MoldSystemDesigner
from core.side_action_designer import SideActionDesigner
from core.mold_cam import MoldCAMDesigner
from core.mold_machining import CollisionDetector, ToolpathOptimizer, EDMElectrodeDesigner, MachiningSimulator
from core.cad_exporter import CADExporter
from services.auth_service import get_current_active_user
from models.database import User

logger = get_logger(__name__)

router = APIRouter()

file_handler = FileHandler()
cavity_layout_optimizer = CavityLayoutOptimizer()
mold_system_designer = MoldSystemDesigner()
side_action_designer = SideActionDesigner()
mold_cam_designer = MoldCAMDesigner()
collision_detector = CollisionDetector()
toolpath_optimizer = ToolpathOptimizer()
edm_designer = EDMElectrodeDesigner()
machining_simulator = MachiningSimulator()
cad_exporter = CADExporter()


@router.get("/health")
@router.post("/health")
async def health():
    task_count = await redis_task_manager.get_task_count()
    return {
        "status": "healthy",
        "pythonocc": True,
        "total_tasks": task_count,
        "redis_connected": redis_task_manager.is_connected,
    }


@router.post("/upload")
async def upload_stp(
    file: UploadFile = File(...),
    material: Optional[str] = "ABS",
    db_session: AsyncSession = Depends(get_db_session),
    current_user: User = Depends(get_current_active_user)
):
    """上传STP文件并存储到数据库"""
    logger.info(
        f"[UPLOAD] 用户={current_user.username}(id={current_user.id}) "
        f"文件={file.filename} 材料={material}"
    )

    if not file.filename.lower().endswith(('.stp', '.step')):
        logger.warning(f"[UPLOAD] 拒绝: 不支持的文件类型 - {file.filename}")
        raise HTTPException(400, "只支持STP/STEP文件")

    task_id = str(uuid.uuid4())

    # 保存文件
    file_path, file_size = await file_handler.save_uploaded_file(file)
    logger.info(f"[UPLOAD] 文件已保存: {file_path} ({file_size} bytes), task_id={task_id}")

    # 创建存储集成服务实例
    storage_service = StorageIntegrationService()

    # 保存STP文件到RustFS + PostgreSQL
    stp_file = await storage_service.save_stp_file(
        session=db_session,
        file_path=file_path,
        original_filename=file.filename,
        user_id=current_user.id
    )
    logger.info(f"[UPLOAD] STP文件已存入RustFS+PG: stp_file.id={stp_file.id}")

    # 创建处理任务记录
    await storage_service.create_processing_task(db_session, task_id, stp_file.id)

    # 创建任务记录（Redis 为主，内存作为兼容回退）
    task_info = create_task_info(
        task_id=task_id,
        status=ProcessingStatus.PROCESSING,
        filename=file.filename,
        file_path=str(file_path),
        file_size=file_size,
        upload_time=str(datetime.now())
    )
    await redis_task_manager.set_task(task_id, task_info)

    # 后台处理走统一编排服务，避免请求会话在后台失效
    process_params = {"material": material}
    if _use_celery:
        process_stp_task.delay(task_id, str(file_path), stp_file.id, process_params)
        logger.info(f"[UPLOAD] Celery 任务已调度: task_id={task_id}")
    else:
        import asyncio
        from services.processing_service import processing_service
        asyncio.create_task(processing_service.process_file_with_storage(
            task_id, str(file_path), stp_file.id, process_params
        ))
        logger.info(f"[UPLOAD] 直接后台处理: task_id={task_id} (celery 未安装)")

    return {
        "task_id": task_id,
        "status": "processing",
        "message": "文件上传成功，开始处理并存储到数据库",
        "file_info": {
            "filename": file.filename,
            "size": file_size,
            "pythonocc_available": True,
            "database_file_id": stp_file.id
        }
    }


@router.get("/status/{task_id}")
@router.post("/status/{task_id}")
async def get_status(task_id: str, db_session: AsyncSession = Depends(get_db_session)):
    """
    获取任务状态

    优先返回内存中的任务信息；
    如果内存中不存在，则从 PostgreSQL + RustFS 组装一个持久化的任务视图，
    结构与内存任务保持尽量一致，便于前端集中展示总结性信息。
    """
    try:
        task_view = await TaskQueryService.get_task_view(db_session, task_id)
        if task_view is None:
            raise HTTPException(404, "任务不存在")
        return task_view
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取任务状态失败: {e}")
        raise HTTPException(500, f"获取任务状态失败: {str(e)}")


@router.get("/debug/tasks")
@router.post("/debug/tasks")
async def debug_tasks():
    """调试接口：查看所有任务"""
    all_tasks = await redis_task_manager.get_all_tasks()
    return {
        "total_tasks": len(all_tasks),
        "tasks": all_tasks,
        "redis_connected": redis_task_manager.is_connected,
    }


@router.get("/history")
@router.post("/history")
async def get_file_history(db_session: AsyncSession = Depends(get_db_session)):
    """获取按文件名分组的文件历史记录（支持多上传）"""
    storage_service = StorageIntegrationService()
    
    file_groups = await storage_service.get_all_file_groups(db_session)
    
    return {
        "total_files": len(file_groups),
        "files": file_groups
    }


@router.get("/history/{filename}")
@router.post("/history/{filename}")
async def get_file_records(filename: str, db_session: AsyncSession = Depends(get_db_session)):
    """获取指定文件名的所有上传记录（支持多上传历史）"""
    import urllib.parse
    decoded_filename = urllib.parse.unquote(filename)
    
    storage_service = StorageIntegrationService()
    file_records = await storage_service.get_file_history_by_filename(
        db_session, 
        decoded_filename
    )
    
    return file_records


@router.get("/result/{task_id}")
@router.post("/result/{task_id}")
async def result_page(request: Request, task_id: str, db_session: AsyncSession = Depends(get_db_session)):
    """结果详情页面"""
    from sqlalchemy import select
    from models.database import ProcessingTask, STPFile, GeometryData, MoldCavityData, HTMLFile
    
    # 从数据库查询任务详情
    result = await db_session.execute(
        select(ProcessingTask, STPFile)
        .join(STPFile, ProcessingTask.stp_file_id == STPFile.id)
        .where(ProcessingTask.task_id == task_id)
    )
    
    task_record = result.first()
    
    if not task_record:
        raise HTTPException(404, "任务不存在")
    
    task, stp_file = task_record
    
    # 构建任务详情数据（先只包含基本数据）
    task_data = {
        "task_id": task.task_id,
        "filename": stp_file.original_filename if stp_file else "",
        "file_size": stp_file.file_size if stp_file else 0,
        "status": task.status,
        "progress": task.progress,
        "current_step": task.current_step,
        "created_at": task.created_time.isoformat() if task.created_time else "",
        "completed_at": task.completed_time.isoformat() if task.completed_time else "",
        "error": task.error_message if task.error_message else ""
    }
    
    from fastapi.templating import Jinja2Templates
    import os
    # 简化路径配置，直接使用当前工作目录下的templates文件夹
    templates_dir = os.path.join(os.getcwd(), "templates")
    templates = Jinja2Templates(directory=templates_dir)
    return templates.TemplateResponse("result.html", {
        "request": request,
        "task": task_data,
        "pythonocc_available": True,
        "version": "3.0.0"
    })


# ==================== P3 新增 API ====================

@router.post("/optimize-layout")
async def optimize_cavity_layout(
    request: Request,
    current_user: User = Depends(get_current_active_user),
):
    """多型腔布局优化"""
    body = await request.json()
    product_bbox = body.get("product_bbox", {"dimensions": [100, 100, 50]})
    cavity_count = body.get("cavity_count", 1)
    mold_base_size = body.get("mold_base_size")
    layout_type = body.get("layout_type", "auto")

    if cavity_count < 1 or cavity_count > 64:
        raise HTTPException(400, "型腔数量必须在 1-64 之间")

    result = cavity_layout_optimizer.optimize_layout(
        product_bbox=product_bbox,
        cavity_count=cavity_count,
        mold_base_size=mold_base_size,
        layout_type=layout_type,
    )

    return {"status": "success", "data": result}


@router.post("/design-cooling")
async def design_cooling_system(
    request: Request,
    current_user: User = Depends(get_current_active_user),
):
    """冷却系统设计"""
    body = await request.json()
    mold_size = body.get("mold_size", {"length": 300, "width": 300, "height": 200})
    product_bbox = body.get("product_bbox", {"dimensions": [100, 100, 50]})
    material = body.get("material", "ABS")
    cavity_count = body.get("cavity_count", 1)
    cycle_time_target = body.get("cycle_time_target")

    from core.mold_system_designer import CoolingSystemDesigner
    designer = CoolingSystemDesigner()
    result = designer.design_cooling_system(
        mold_size=mold_size,
        product_bbox=product_bbox,
        material=material,
        cavity_count=cavity_count,
        cycle_time_target=cycle_time_target,
    )

    return {"status": "success", "data": result}


@router.post("/design-gating")
async def design_gating_system(
    request: Request,
    current_user: User = Depends(get_current_active_user),
):
    """浇注系统设计"""
    body = await request.json()
    product_bbox = body.get("product_bbox", {"dimensions": [100, 100, 50]})
    material = body.get("material", "ABS")
    cavity_count = body.get("cavity_count", 1)
    gate_type = body.get("gate_type", "auto")
    layout_positions = body.get("layout_positions")

    from core.mold_system_designer import GatingSystemDesigner
    designer = GatingSystemDesigner()
    result = designer.design_gating_system(
        product_bbox=product_bbox,
        material=material,
        cavity_count=cavity_count,
        gate_type=gate_type,
        layout_positions=layout_positions,
    )

    return {"status": "success", "data": result}


@router.post("/design-mold-system")
async def design_complete_mold_system(
    request: Request,
    current_user: User = Depends(get_current_active_user),
):
    """综合模具系统设计（冷却+浇注）"""
    body = await request.json()
    mold_size = body.get("mold_size", {"length": 300, "width": 300, "height": 200})
    product_bbox = body.get("product_bbox", {"dimensions": [100, 100, 50]})
    material = body.get("material", "ABS")
    cavity_count = body.get("cavity_count", 1)
    gate_type = body.get("gate_type", "auto")
    cycle_time_target = body.get("cycle_time_target")
    layout_positions = body.get("layout_positions")

    result = mold_system_designer.design_complete_system(
        mold_size=mold_size,
        product_bbox=product_bbox,
        material=material,
        cavity_count=cavity_count,
        gate_type=gate_type,
        cycle_time_target=cycle_time_target,
        layout_positions=layout_positions,
    )

    return {"status": "success", "data": result}


@router.post("/ai-parting-detect")
async def ai_parting_surface_detect(
    request: Request,
    current_user: User = Depends(get_current_active_user),
):
    """AI 分型面检测"""
    body = await request.json()
    task_id = body.get("task_id")

    if not task_id:
        raise HTTPException(404, "缺少 task_id")

    task_data = await redis_task_manager.get_task(task_id)
    if not task_data:
        raise HTTPException(404, "任务不存在")

    geometry_data = task_data.get("geometry_data")
    if not geometry_data:
        raise HTTPException(400, "该任务尚未完成几何分析")

    from core.ai_parting_detector import AIPartingSurfaceDetectorV2
    detector = AIPartingSurfaceDetectorV2(use_gnn=True)

    result = detector._detect_with_geometry(None, geometry_data)

    return {"status": "success", "data": result}


@router.post("/detect-undercuts")
async def detect_undercuts(
    request: Request,
    current_user: User = Depends(get_current_active_user),
):
    """倒扣区域检测与滑块/斜顶机构设计"""
    body = await request.json()
    task_id = body.get("task_id")
    parting_direction = body.get("parting_direction", [0, 0, 1])
    mold_size = body.get("mold_size", {"length": 300, "width": 300, "height": 200})

    if not task_id:
        raise HTTPException(404, "缺少 task_id")

    task_data = await redis_task_manager.get_task(task_id)
    if not task_data:
        raise HTTPException(404, "任务不存在")

    geometry_data = task_data.get("geometry_data")
    if not geometry_data:
        raise HTTPException(400, "该任务尚未完成几何分析")

    result = side_action_designer.analyze_and_design(
        shape=None, parting_direction=parting_direction, mold_size=mold_size
    )

    return {"status": "success", "data": result}


@router.post("/design-cam")
async def design_mold_cam(
    request: Request,
    current_user: User = Depends(get_current_active_user),
):
    """模具CAM刀路设计"""
    body = await request.json()
    cavity_bbox = body.get("cavity_bbox", {"dimensions": [100, 100, 50], "min": [-50, -50, -25], "max": [50, 50, 25]})
    stock_bbox = body.get("stock_bbox", {"dimensions": [150, 150, 100], "min": [-75, -75, -50], "max": [75, 75, 50]})
    mold_steel = body.get("mold_steel", "P20")
    surface_quality = body.get("surface_quality", "standard")
    controller = body.get("controller", "fanuc")

    result = mold_cam_designer.design_mold_cam(
        cavity_bbox=cavity_bbox,
        stock_bbox=stock_bbox,
        mold_steel=mold_steel,
        surface_quality=surface_quality,
        controller=controller,
    )

    return {"status": "success", "data": result}


@router.post("/check-collision")
async def check_toolpath_collision(
    request: Request,
    current_user: User = Depends(get_current_active_user),
):
    """刀路碰撞检测"""
    body = await request.json()
    toolpath_points = body.get("toolpath_points", [[0, 0, 50], [10, 10, -5], [20, 20, -10]])
    tool = body.get("tool", {"diameter": 10, "flute_length": 30, "shank_diameter": 10})
    stock_bbox = body.get("stock_bbox", {"min": [-50, -50, -25], "max": [50, 50, 25]})
    clamp_positions = body.get("clamp_positions")

    result = collision_detector.check_toolpath_safety(
        toolpath_points, tool, stock_bbox, clamp_positions
    )

    return {"status": "success", "data": result}


@router.post("/optimize-toolpath")
async def optimize_toolpath(
    request: Request,
    current_user: User = Depends(get_current_active_user),
):
    """刀路优化"""
    body = await request.json()
    toolpath_points = body.get("toolpath_points", [[0, 0, 50], [10, 10, -5], [20, 20, -10]])
    cutting_params = body.get("cutting_params", {"feed_rate_mm_min": 500})
    stock_bbox = body.get("stock_bbox")

    result = toolpath_optimizer.optimize_toolpath(
        toolpath_points, cutting_params, stock_bbox
    )

    return {"status": "success", "data": result}


@router.post("/design-electrodes")
async def design_edm_electrodes(
    request: Request,
    current_user: User = Depends(get_current_active_user),
):
    """EDM电极设计"""
    body = await request.json()
    undercut_regions = body.get("undercut_regions", [{"center": [0, 0, 0], "area": 100, "type": "undercut"}])
    cavity_bbox = body.get("cavity_bbox", {"dimensions": [100, 100, 50]})
    material = body.get("material", "copper")
    spark_gap = body.get("spark_gap", 0.05)
    overburn = body.get("overburn", 0.1)

    result = edm_designer.design_electrodes(
        undercut_regions, cavity_bbox, material, spark_gap, overburn
    )

    return {"status": "success", "data": result}


@router.post("/simulate-machining")
async def simulate_machining(
    request: Request,
    current_user: User = Depends(get_current_active_user),
):
    """加工仿真"""
    body = await request.json()
    operations = body.get("operations", [{"strategy": "z_level_roughing", "levels": [{"z": -5}]}])
    stock_bbox = body.get("stock_bbox", {"dimensions": [100, 100, 50], "min": [-50, -50, -25], "max": [50, 50, 25]})
    resolution = body.get("resolution", 2.0)

    result = machining_simulator.simulate_machining(
        operations, stock_bbox, resolution
    )

    return {"status": "success", "data": result}


# ==================== CAD 导出 API ====================

@router.post("/export-mold")
async def export_mold_results(
    request: Request,
    current_user: User = Depends(get_current_active_user),
):
    """导出模具设计结果（STEP/IGES/STL/BRep）"""
    body = await request.json()
    task_id = body.get("task_id")
    formats = body.get("formats", ["step", "stl"])
    components = body.get("components", ["cavity", "core"])
    scheme_id = body.get("scheme_id")

    if not task_id:
        raise HTTPException(404, "缺少 task_id")

    cavity_shapes = processing_service.get_export_shapes(task_id, scheme_id=scheme_id)
    if not cavity_shapes:
        task_data = await redis_task_manager.get_task(task_id)
        if not task_data:
            raise HTTPException(404, "任务不存在")
        filename = task_data.get("filename", f"mold_{task_id}")
        raise HTTPException(
            400,
            f"该任务的 OCC 形状数据已过期（仅保留 STEP 导出文件），"
            f"请通过历史页面的下载链接获取已导出的 STEP 文件",
        )

    task_data = await redis_task_manager.get_task(task_id)
    base_filename = (
        Path(task_data.get("filename", f"mold_{task_id}")).stem
        if task_data
        else f"mold_{task_id}"
    )

    result = cad_exporter.export_mold_results(
        cavity_data=cavity_shapes,
        base_filename=base_filename,
        formats=formats,
        components=components,
    )

    return {"status": "success", "data": result}


@router.get("/export-download/{filepath:path}")
async def download_export_file(
    filepath: str,
    current_user: User = Depends(get_current_active_user),
):
    """下载导出的CAD文件"""
    from fastapi.responses import FileResponse

    full_path = os.path.join(cad_exporter.output_dir, filepath)

    if not os.path.exists(full_path):
        raise HTTPException(404, "文件不存在")

    if not os.path.abspath(full_path).startswith(os.path.abspath(cad_exporter.output_dir)):
        raise HTTPException(403, "禁止访问")

    media_types = {
        ".step": "application/step",
        ".stp": "application/step",
        ".iges": "application/iges",
        ".igs": "application/iges",
        ".stl": "model/stl",
        ".brep": "application/octet-stream",
    }

    ext = Path(full_path).suffix.lower()
    media_type = media_types.get(ext, "application/octet-stream")

    return FileResponse(
        full_path,
        media_type=media_type,
        filename=os.path.basename(full_path),
    )


@router.get("/export-recommendations")
async def get_export_recommendations(
    target: str = "ug",
    current_user: User = Depends(get_current_active_user),
):
    """获取导出格式建议（UG/FreeCAD/SolidWorks）"""
    result = cad_exporter.get_export_recommendations(target)
    return {"status": "success", "data": result}