geMoldInsight/src/api/routes.py

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

from models.schemas import ProcessingStatus, create_task_info
from core.stp_parser import STPParser
from core.geometry_analyzer import GeometryAnalyzer
from utils.file_handler import FileHandler
from utils.html_generator import HTMLGenerator
from services.storage_integration_rustfs import StorageIntegrationService
from database.database import get_db_session
from utils.logger import get_logger
from sqlalchemy.ext.asyncio import AsyncSession
from core.mold_generator import MoldCavityGenerator
from core.mesh_generator import MeshGenerator

logger = get_logger(__name__)

router = APIRouter()

stp_parser = STPParser()
geometry_analyzer = GeometryAnalyzer()
file_handler = FileHandler()
html_generator = HTMLGenerator()
mold_generator = MoldCavityGenerator(shrinkage_rate=0.005)
mesh_generator = MeshGenerator(quality="medium")

tasks = {}


@router.get("/health")
@router.post("/health")
async def health():
    return {
        "status": "healthy",
        "pythonocc": True,
        "total_tasks": len(tasks)
    }


@router.post("/upload")
async def upload_stp(
    background_tasks: BackgroundTasks, 
    file: UploadFile = File(...),
    db_session: AsyncSession = Depends(get_db_session)
):
    """上传STP文件并存储到数据库"""

    if not file.filename.lower().endswith(('.stp', '.step')):
        raise HTTPException(400, "只支持STP/STEP文件")

    task_id = str(uuid.uuid4())

    # 保存文件
    file_path = await file_handler.save_uploaded_file(file)
    content = await file.read()

    # 创建存储集成服务实例
    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
    )

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

    # 创建内存任务记录
    tasks[task_id] = create_task_info(
        task_id=task_id,
        status=ProcessingStatus.PROCESSING,
        filename=file.filename,
        file_path=str(file_path),
        file_size=len(content),
        upload_time=str(datetime.now())
    )

    # 后台处理（包含数据库存储）
    background_tasks.add_task(process_file_with_storage, task_id, file_path, stp_file.id, db_session)

    return {
        "task_id": task_id,
        "status": "processing",
        "message": "文件上传成功，开始处理并存储到数据库",
        "file_info": {
            "filename": file.filename,
            "size": len(content),
            "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 组装一个持久化的任务视图，
    结构与内存任务保持尽量一致，便于前端集中展示总结性信息。
    """
    # 1. 内存任务（进行中的任务）
    if task_id in tasks:
        task = tasks[task_id]
        logger.info(f"返回内存任务状态: {task_id} - {task['status']}")
        return task

    # 2. 持久化任务（已完成/失败，或服务重启后的任务）
    from sqlalchemy import select
    from models.database import ProcessingTask, STPFile, GeometryData, MeshData, MoldCavityData
    storage_service = StorageIntegrationService()

    # 查询任务和文件元数据
    result = await db_session.execute(
        select(ProcessingTask, STPFile)
        .join(STPFile, ProcessingTask.stp_file_id == STPFile.id)
        .where(ProcessingTask.task_id == task_id)
    )
    row = result.first()
    if not row:
        raise HTTPException(404, "任务不存在")

    processing_task, stp_file = row

    # 从 RustFS 取几何 / 型腔 / 网格详细 JSON（小量数据，便于前端展示汇总）
    file_with_data = await storage_service.get_stp_file_with_data(
        db_session, stp_file_id=stp_file.id
    )

    geometry_json: Optional[Dict[str, Any]] = None
    if file_with_data.get("geometry_data"):
        # save_geometry_data 保存时可能有两种格式：
        # 1. 直接保存 geometry_data 字典
        # 2. 保存 {"geometry_data": {...}} 格式
        # get_stp_file_with_data 返回的是从RustFS下载的原始JSON
        geo_raw = file_with_data["geometry_data"]
        if isinstance(geo_raw, dict):
            # 如果是包装格式，提取内部数据
            if "geometry_data" in geo_raw:
                geometry_json = geo_raw["geometry_data"]
            else:
                # 直接就是几何数据字典
                geometry_json = geo_raw

    cavity_json: Optional[Dict[str, Any]] = file_with_data.get("mold_cavity_data")

    # 组装网格摘要：从 MeshData 表中读出摘要字段，避免把完整网格 JSON 丢给前端
    mesh_summary = None
    mesh_record = await db_session.execute(
        select(MeshData).where(MeshData.stp_file_id == stp_file.id)
    )
    mesh_record = mesh_record.scalar_one_or_none()
    if mesh_record:
        mesh_summary = {
            "vertex_count": mesh_record.vertex_count,
            "face_count": mesh_record.face_count,
            "point_count": mesh_record.point_count,
            "quality": mesh_record.quality,
        }

    # 构造与内存任务兼容的任务视图
    task_view = {
        "task_id": processing_task.task_id,
        "status": processing_task.status,
        "filename": stp_file.original_filename if stp_file else "",
        "file_path": stp_file.file_path or "",
        "file_size": stp_file.file_size if stp_file else 0,
        "upload_time": processing_task.created_time.isoformat()
        if processing_task.created_time
        else "",
        "completed_at": processing_task.completed_time.isoformat()
        if processing_task.completed_time
        else "",
        "geometry_data": geometry_json,
        # key_info 沿用之前的结构，直接使用详细型腔 JSON，前端已按该结构解析
        "key_info": cavity_json,
        "mesh_summary": mesh_summary,
        "analysis_result": None,
        "error": processing_task.error_message or stp_file.error_message or None,
    }

    logger.info(f"返回持久化任务状态: {task_id} - {processing_task.status}")
    return task_view


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


@router.get("/api/history")
@router.post("/api/history")
async def get_file_history(db_session: AsyncSession = Depends(get_db_session)):
    """获取按文件名分组的文件历史记录"""
    from sqlalchemy import select
    from models.database import ProcessingTask, STPFile
    
    # 从数据库查询所有处理任务
    result = await db_session.execute(
        select(ProcessingTask, STPFile)
        .join(STPFile, ProcessingTask.stp_file_id == STPFile.id)
        .order_by(ProcessingTask.created_time.desc())
    )
    
    tasks = result.all()
    
    # 按文件名分组
    file_groups = {}
    for task, stp_file in tasks:
        filename = stp_file.original_filename
        if filename not in file_groups:
            file_groups[filename] = []
        
        file_groups[filename].append({
            "task_id": task.task_id,
            "filename": filename,
            "upload_time": task.created_time.isoformat() if task.created_time else "",
            "status": task.status,
            "file_size": stp_file.file_size
        })
    
    # 构建返回数据
    files = []
    for filename, file_tasks in file_groups.items():
        # 按上传时间排序
        file_tasks.sort(key=lambda x: x.get("upload_time", ""), reverse=True)
        
        files.append({
            "filename": filename,
            "record_count": len(file_tasks),
            "last_upload": file_tasks[0].get("upload_time", ""),
            "first_upload": file_tasks[-1].get("upload_time", "") if len(file_tasks) > 1 else ""
        })
    
    # 按最后上传时间排序
    files.sort(key=lambda x: x["last_upload"], reverse=True)
    
    return {
        "total_files": len(files),
        "files": files
    }


@router.get("/api/history/{filename}")
@router.post("/api/history/{filename}")
async def get_file_records(filename: str, db_session: AsyncSession = Depends(get_db_session)):
    """获取指定文件名的所有记录"""
    # URL解码文件名
    import urllib.parse
    decoded_filename = urllib.parse.unquote(filename)
    
    # 从数据库查询指定文件名的所有处理任务
    from sqlalchemy import select
    from models.database import ProcessingTask, STPFile
    
    result = await db_session.execute(
        select(ProcessingTask, STPFile)
        .join(STPFile, ProcessingTask.stp_file_id == STPFile.id)
        .where(STPFile.original_filename == decoded_filename)
        .order_by(ProcessingTask.created_time.desc())
    )
    
    tasks = result.all()
    
    # 构建返回数据
    file_records = []
    for task, stp_file in tasks:
        file_records.append({
            "task_id": task.task_id,
            "filename": stp_file.original_filename,
            "file_size": stp_file.file_size,
            "upload_time": task.created_time.isoformat() if task.created_time else "",
            "status": task.status,
            "completed_at": task.completed_time.isoformat() if task.completed_time else ""
        })
    
    # 按上传时间排序（最新的在前）
    file_records.sort(key=lambda x: x.get("upload_time", ""), reverse=True)
    
    return file_records


@router.get("/history")
@router.post("/history")
async def history_page(request: Request):
    """历史记录页面"""
    from fastapi.templating import Jinja2Templates
    import os
    # 简化路径配置，直接使用当前工作目录下的templates文件夹
    templates_dir = os.path.join(os.getcwd(), "templates")
    templates = Jinja2Templates(directory=templates_dir)
    return templates.TemplateResponse("history.html", {
        "request": request,
        "pythonocc_available": True,
        "version": "3.0.0"
    })


@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"
    })


async def process_file_with_storage(
        task_id: str,
        file_path: str,
        stp_file_id: int,
        db_session: AsyncSession
):
    """处理文件的后台任务"""

    storage_service = StorageIntegrationService()

    try:
        logger.info(f"开始处理文件并生成模具型腔: {file_path}")
        
        # 设置处理超时（5分钟）
        import asyncio
        timeout_seconds = 300  # 5分钟
        
        async def process_with_timeout():
            # 处理逻辑将在下面添加
            pass
            
        # 使用超时保护
        try:
            await asyncio.wait_for(process_file_core(storage_service, task_id, file_path, stp_file_id, db_session), timeout_seconds)
        except asyncio.TimeoutError:
            logger.error(f"处理超时: {task_id}")
            raise Exception(f"处理超时，超过{timeout_seconds}秒未完成")
            
    except Exception as e:
        logger.error(f"模具型腔生成失败: {e}")

        await storage_service.update_stp_file_status(db_session, stp_file_id, "failed")
        await storage_service.update_task_status(
            db_session, task_id, "failed", error_message=str(e)
        )

        tasks[task_id]["status"] = ProcessingStatus.FAILED
        tasks[task_id]["error"] = str(e)
        tasks[task_id]["completed_at"] = str(datetime.now())
        return


async def process_file_core(
        storage_service: StorageIntegrationService,
        task_id: str,
        file_path: str,
        stp_file_id: int,
        db_session: AsyncSession
):
    """核心处理逻辑"""
    
    try:
        logger.info(f"开始处理文件并生成模具型腔: {file_path}")

        # 更新任务状态
        await storage_service.update_task_status(
            db_session, task_id, "processing", 20, "解析STP文件"
        )

        # 1. 解析STP文件
        await storage_service.update_task_status(
            db_session, task_id, "processing", 20, "解析STP文件"
        )
        
        # 使用STPParser类进行真实解析
        shape = stp_parser.load_step_file(Path(file_path))
        geometry_data = stp_parser.analyze_geometry(shape)

        # 2. 生成网格数据并持久化（详细 JSON 存 RustFS，摘要写 PostgreSQL）
        await storage_service.update_task_status(
            db_session, task_id, "processing", 30, "生成网格数据"
        )

        mesh_result = None
        mesh_json = None
        try:
            mesh_result = mesh_generator.generate_mesh_from_shape(shape)

            tri_mesh = mesh_result.get("trimesh_mesh")
            pointcloud = mesh_result.get("pointcloud") or {}

            if tri_mesh is not None:
                # 顶点和面转为 JSON 可序列化
                vertices = tri_mesh.vertices.tolist()
                faces = tri_mesh.faces.tolist()

                # 使用已计算的几何边界框，避免重复计算
                bbox = geometry_data.get("bounding_box", {})

                mesh_json = {
                    "metadata": {
                        "file_name": Path(file_path).name,
                        "generated_at": datetime.now().isoformat(),
                        "quality": "medium",
                        "vertex_count": len(vertices),
                        "face_count": len(faces),
                        "point_count": pointcloud.get("count"),
                    },
                    "mesh": {
                        "vertices": vertices,
                        "faces": faces,
                    },
                    "pointcloud": pointcloud,
                    "bounding_box": bbox,
                }

                await storage_service.save_mesh_data(
                    db_session,
                    stp_file_id=stp_file_id,
                    mesh_json=mesh_json,
                    quality="medium",
                )
                # 将简要网格摘要写入内存任务，便于前端展示汇总信息
                tasks[task_id]["mesh_summary"] = {
                    "vertex_count": len(vertices),
                    "face_count": len(faces),
                    "point_count": pointcloud.get("count"),
                    "quality": "medium",
                }
        except Exception as mesh_err:
            # 网格失败不影响整体流程，只记录日志
            logger.warning(f"网格生成或保存失败，不影响主流程: {mesh_err}")

        # 3. 生成模具型腔（模拟数据）
        await storage_service.update_task_status(
            db_session, task_id, "processing", 40, "生成模具型腔（模拟）"
        )

        cavity_data = {
            "cavity_count": 1,
            "cavity_dimensions": {"length": 100, "width": 80, "height": 50},
            "runner_system": "cold_runner",
            "gating_type": "edge_gate"
        }

        # 4. 生成详细JSON数据（使用计算值）
        await storage_service.update_task_status(
            db_session, task_id, "processing", 60, "生成型腔详细数据"
        )

        # 使用模具生成器计算各项参数
        volume_mm3 = geometry_data.get("volume", 0)
        surface_area_mm2 = geometry_data.get("surface_area", 0)
        bbox = geometry_data.get("bounding_box", {})
        bbox_dims = bbox.get("dimensions", [0, 0, 0])

        # 计算产品重量（ABS密度：1.05 g/cm³）
        volume_cm3 = volume_mm3 / 1000
        product_weight_g = volume_cm3 * 1.05

        # 计算投影面积（取X、Y方向）
        if len(bbox_dims) >= 2:
            projected_area_cm2 = (bbox_dims[0] * bbox_dims[1]) / 100
        else:
            projected_area_cm2 = 0

        # 计算夹紧力（投影面积 × 注塑压力600 kg/cm²，转换为吨）
        clamping_force_ton = int(projected_area_cm2 * 600 / 1000)

        # 计算壁厚范围
        if surface_area_mm2 > 0 and volume_mm3 > 0:
            avg_thickness_mm = (volume_mm3 / surface_area_mm2) * 0.6
            wall_thickness_min = avg_thickness_mm * 0.7
            wall_thickness_max = avg_thickness_mm * 1.3
        else:
            avg_thickness_mm = 2.5
            wall_thickness_min = 2.0
            wall_thickness_max = 3.0

        # 计算复杂度评分
        if surface_area_mm2 > 0 and volume_mm3 > 0:
            complexity_score = min((avg_thickness_mm / 5.0), 1.0)
        else:
            complexity_score = 0.5

        # 计算模具尺寸（基于产品尺寸 + 模具边距）
        mold_length = max(bbox_dims[0] if len(bbox_dims) > 0 else 120, 120) + 40
        mold_width = max(bbox_dims[1] if len(bbox_dims) > 1 else 100, 100) + 40
        mold_height = max(bbox_dims[2] if len(bbox_dims) > 2 else 60, 60) + 50

        # 计算分型线长度（基于产品周长）
        if len(bbox_dims) >= 2:
            parting_line_length = 2 * (bbox_dims[0] + bbox_dims[1])
        else:
            parting_line_length = 0

        # 估算成型周期（基于体积）
        # 周期 = 冷却时间 + 注塑时间 + 开合模时间
        cooling_time = (wall_thickness_max ** 2) * 5  # 简化公式
        injection_time = max(5, volume_cm3 / 50)  # 注塑时间
        cycle_time = cooling_time + injection_time + 8  # 开合模约8秒

        detailed_cavity_json = {
            "metadata": {
                "file_name": Path(file_path).name,
                "analysis_date": datetime.now().isoformat(),
                "shrinkage_rate": 0.005,
                "draft_angle": 2.0
            },
            "product_analysis": {
                "volume": volume_mm3,
                "surface_area": surface_area_mm2,
                "bounding_box": bbox
            },
            "manufacturing_info": {
                "recommended_material": "ABS",
                "estimated_clamping_force": f"{clamping_force_ton} 吨",
                "estimated_mold_size": {
                    "length": int(mold_length),
                    "width": int(mold_width),
                    "height": int(mold_height)
                },
                "mold_material": "铝合金7075",
                "mold_hardness": "HB 150-170",
                "surface_finish": "Ra 0.8 μm",
                "parting_line_length": f"{parting_line_length:.2f} mm",
                "estimated_cycle_time": f"{int(cycle_time)} 秒"
            },
            "mold_cavities": {
                "cavity_count": 1,
                "cavity_key_info": {
                    "geometric_characteristics": {
                        "product_weight": f"{product_weight_g:.2f} g",
                        "wall_thickness_range": f"{wall_thickness_min:.2f} - {wall_thickness_max:.2f} mm",
                        "complexity_score": round(complexity_score, 2),
                        "product_volume": f"{volume_cm3:.2f} cm³"
                    },
                    "quality_considerations": {
                        "potential_weld_lines": "center",
                        "sink_mark_areas": "thick_sections",
                        "warpage_risk": "low"
                    }
                }
            }
        }

        # 5. 生成关键信息（模拟）
        cavity_key_info = detailed_cavity_json["mold_cavities"]["cavity_key_info"]

        # 6. 保存几何数据到数据库
        await storage_service.update_task_status(
            db_session, task_id, "processing", 70, "保存几何数据"
        )

        geometry_record = await storage_service.save_geometry_data(
            db_session,
            stp_file_id,
            geometry_data,
            geometry_data.get("analysis_method", "mold_cavity")
        )

        # 7. 保存模具型腔数据
        await storage_service.save_mold_cavity_data(
            db_session,
            stp_file_id,
            detailed_cavity_json
        )

        # 8. 生成HTML可视化（包含型腔信息）
        await storage_service.update_task_status(
            db_session, task_id, "processing", 85, "生成可视化报告"
        )

        html_file_path = html_generator.generate_and_save_visualization(
            geometry_data,
            Path(file_path).name,
            cavity_data=detailed_cavity_json
        )

        # 保存HTML文件信息
        html_record = await storage_service.save_html_file(
            db_session,
            stp_file_id,
            Path(html_file_path).name,
            html_file_path
        )

        # 9. 分析模具设计
        analysis_result = geometry_analyzer.analyze_mold_design(geometry_data)

        # 10. 完成处理
        await storage_service.update_stp_file_status(db_session, stp_file_id, "completed")
        await storage_service.update_task_status(
            db_session, task_id, "completed", 100, "模具型腔生成完成"
        )

        # 更新内存任务状态
        tasks[task_id]["geometry_data"] = geometry_data
        tasks[task_id]["analysis_result"] = analysis_result
        tasks[task_id]["cavity_data"] = detailed_cavity_json
        tasks[task_id]["key_info"] = detailed_cavity_json  # 传递完整数据给前端
        tasks[task_id]["status"] = ProcessingStatus.COMPLETED
        tasks[task_id]["completed_at"] = str(datetime.now())

        # 调试日志
        logger.info(f"模具型腔生成完成: {task_id}")
        logger.info(f"key_info metadata: {detailed_cavity_json.get('metadata', {})}")
        logger.info(f"key_info manufacturing_info: {detailed_cavity_json.get('manufacturing_info', {})}")
        logger.info(f"key_info geometric_characteristics: {detailed_cavity_json.get('mold_cavities', {}).get('cavity_key_info', {}).get('geometric_characteristics', {})}")

    except Exception as e:
        logger.error(f"模具型腔生成失败: {e}")

        await storage_service.update_stp_file_status(db_session, stp_file_id, "failed")
        await storage_service.update_task_status(
            db_session, task_id, "failed", error_message=str(e)
        )

        tasks[task_id]["status"] = ProcessingStatus.FAILED
        tasks[task_id]["error"] = str(e)
        tasks[task_id]["completed_at"] = str(datetime.now())