geMoldInsight/src/api/routes.py

# api/routes.py
from fastapi import APIRouter, UploadFile, File, HTTPException, BackgroundTasks, Request, Depends
from typing import Optional
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

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)  # ABS材料

# 内存中的任务存储
tasks = {}


@router.get("/")
async def read_root(request: Request):
    """主页面"""
    from fastapi.templating import Jinja2Templates
    import os
    # 使用绝对路径确保模板目录正确
    templates_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "templates")
    templates = Jinja2Templates(directory=templates_dir)
    return templates.TemplateResponse("index.html", {
        "request": request,
        "pythonocc_available": True,
        "version": "3.0.0"
    })


@router.get("/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}")
async def get_status(task_id: str):
    """获取任务状态"""
    if task_id not in tasks:
        raise HTTPException(404, "任务不存在")

    task = tasks[task_id]
    logger.info(f"返回任务状态: {task_id} - {task['status']}")
    return task


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


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. 生成模具型腔（模拟数据）
        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"
        }

        # 3. 生成详细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": "Aluminum Alloy 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"
                    }
                }
            }
        }

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

        # 5. 保存几何数据到数据库
        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")
        )

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

        # 7. 生成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
        )

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

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

        # 9. 完成处理
        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"] = cavity_key_info
        tasks[task_id]["status"] = ProcessingStatus.COMPLETED
        tasks[task_id]["completed_at"] = str(datetime.now())

        logger.info(f"模具型腔生成完成: {task_id}")

    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())