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, "生成型腔详细数据（模拟）"
        )

        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": geometry_data["volume"],
                "surface_area": geometry_data["surface_area"],
                "bounding_box": geometry_data["bounding_box"]
            },
            "manufacturing_info": {
                "recommended_material": "ABS",
                "estimated_clamping_force": "150 吨",
                "estimated_mold_size": {
                    "length": 120,
                    "width": 100,
                    "height": 60
                }
            },
            "mold_cavities": {
                "cavity_count": 1,
                "cavity_key_info": {
                    "geometric_characteristics": {
                        "product_weight": "1.2 g",
                        "wall_thickness_range": "1.5-3.0 mm",
                        "complexity_score": 0.7
                    },
                    "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())