geMoldInsight/src/api/v1/advanced_router.py

from pathlib import Path
import os
from datetime import datetime
from urllib.parse import quote

from fastapi import APIRouter, Depends, HTTPException, Request
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from services.auth_service import get_current_active_user
from services.redis_task_manager import redis_task_manager
from services.processing_service import processing_service
from services.storage_integration_rustfs import StorageIntegrationService
from services.task_query_service import TaskQueryService
from database.database import get_db_session
from models.database import User
from models.database import ProcessingTask, STPFile
from core.cad_exporter import CADExporter
from utils.logger import get_logger

logger = get_logger(__name__)

router = APIRouter()
cad_exporter = CADExporter()
storage_service = StorageIntegrationService()

_api_routes_cache = {}


def _get_cached(key):
    global _api_routes_cache
    if key not in _api_routes_cache:
        try:
            from api import routes
        except Exception as e:
            logger.warning(f"api.routes 模块加载失败: {e}")
            _api_routes_cache["__error__"] = str(e)
            return None
        _api_routes_cache.clear()
        _api_routes_cache.update({
            "tasks": routes.tasks,
            "cavity_layout_optimizer": routes.cavity_layout_optimizer,
            "mold_system_designer": routes.mold_system_designer,
            "side_action_designer": routes.side_action_designer,
            "mold_cam_designer": routes.mold_cam_designer,
            "collision_detector": routes.collision_detector,
            "toolpath_optimizer": routes.toolpath_optimizer,
            "edm_designer": routes.edm_designer,
            "machining_simulator": routes.machining_simulator,
            "cad_exporter": routes.cad_exporter,
        })
    return _api_routes_cache.get(key)


async def _get_task_data(task_id: str) -> dict:
    task = await redis_task_manager.get_task(task_id)
    if task:
        return task
    tasks = _get_cached("tasks")
    if tasks and task_id in tasks:
        return tasks[task_id]
    return None


async def _ensure_task_access(
    db_session: AsyncSession,
    task_id: str,
    user_id: int,
):
    row = await db_session.execute(
        select(ProcessingTask, STPFile)
        .join(STPFile, ProcessingTask.stp_file_id == STPFile.id)
        .where(ProcessingTask.task_id == task_id)
    )
    row = row.first()
    if not row:
        raise HTTPException(404, "任务不存在")

    _, stp_file = row
    owner_id = getattr(stp_file, "user_id", None)
    if owner_id is not None and owner_id != user_id:
        raise HTTPException(403, "无权访问该任务的导出文件")

    return row


def _get_export_artifacts(task_data: dict) -> dict:
    if not isinstance(task_data, dict):
        return {}
    direct = task_data.get("export_artifacts")
    if isinstance(direct, dict):
        return direct
    parameters = task_data.get("parameters")
    if isinstance(parameters, dict) and isinstance(parameters.get("export_artifacts"), dict):
        return parameters.get("export_artifacts")
    return {}


def _expand_components(components):
    requested = components or ["cavity", "core"]
    if "all" in requested:
        return ["cavity", "core", "parting_surface"]
    return list(dict.fromkeys(requested))


def _augment_export_files(task_id: str, files):
    items = []
    for file in files or []:
        item = dict(file)
        relative_path = item.get("relative_path")
        if not relative_path and item.get("filepath"):
            relative_path = cad_exporter.get_relative_path(item["filepath"])
        if relative_path:
            relative_path = str(relative_path).replace("\\", "/").strip("/")
            item["relative_path"] = relative_path
            item["download_path"] = f"/api/export-download/{quote(relative_path, safe='/')}?task_id={task_id}"
        items.append(item)
    return items


def _merge_export_artifacts(existing: dict, export_result: dict) -> dict:
    merged = dict(existing or {})
    schemes = dict(merged.get("schemes") or {})
    scheme_id = export_result.get("scheme_id") or "default"
    previous = dict(schemes.get(scheme_id) or {})

    file_map = {}
    for file in previous.get("files", []):
        file_map[(file.get("component"), file.get("format"))] = file
    for file in export_result.get("files", []):
        file_map[(file.get("component"), file.get("format"))] = file

    schemes[scheme_id] = {
        "base_filename": export_result.get("base_filename") or previous.get("base_filename"),
        "generated_at": datetime.now().isoformat(),
        "files": sorted(
            file_map.values(),
            key=lambda item: (item.get("component", ""), item.get("format", "")),
        ),
        "errors": export_result.get("errors", []),
        "total_files": len(file_map),
        "total_errors": len(export_result.get("errors", [])),
    }

    merged["version"] = 1
    merged["task_id"] = export_result.get("task_id") or merged.get("task_id")
    merged["generated_at"] = merged.get("generated_at") or datetime.now().isoformat()
    merged["schemes"] = schemes
    return merged


def _select_persisted_files(task_id: str, task_data: dict, scheme_id: str, formats, components):
    artifacts = _get_export_artifacts(task_data)
    scheme_data = (artifacts.get("schemes") or {}).get(scheme_id)
    if not scheme_data:
        return None

    component_list = _expand_components(components)
    format_list = list(dict.fromkeys(formats or ["step", "stl"]))
    expected = {(component, fmt) for component in component_list for fmt in format_list}

    available = []
    available_keys = set()
    for file in scheme_data.get("files", []):
        component = file.get("component")
        fmt = file.get("format")
        if component not in component_list or fmt not in format_list:
            continue
        relative_path = str(file.get("relative_path") or "").replace("\\", "/").strip("/")
        if not relative_path:
            continue
        full_path = os.path.join(cad_exporter.output_dir, relative_path.replace("/", os.sep))
        if not os.path.exists(full_path):
            continue
        available.append(file)
        available_keys.add((component, fmt))

    if expected and not expected.issubset(available_keys):
        return None

    return _augment_export_files(task_id, available)


@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 之间")
    optimizer = _get_cached("cavity_layout_optimizer")
    if not optimizer:
        raise HTTPException(503, "服务不可用：核心模块未加载")
    result = 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")
    ds = _get_cached("mold_system_designer")
    if not ds:
        raise HTTPException(503, "服务不可用：核心模块未加载")
    result = ds.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),
):
    body = await request.json()
    task_id = body.get("task_id")
    if not task_id:
        raise HTTPException(404, "缺少 task_id")
    task_data = await _get_task_data(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 _get_task_data(task_id)
    if not task_data:
        raise HTTPException(404, "任务不存在")
    sd = _get_cached("side_action_designer")
    if not sd:
        raise HTTPException(503, "服务不可用：核心模块未加载")
    result = sd.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),
):
    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")
    cam = _get_cached("mold_cam_designer")
    if not cam:
        raise HTTPException(503, "服务不可用：核心模块未加载")
    result = cam.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")
    cd = _get_cached("collision_detector")
    if not cd:
        raise HTTPException(503, "服务不可用：核心模块未加载")
    result = cd.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")
    to = _get_cached("toolpath_optimizer")
    if not to:
        raise HTTPException(503, "服务不可用：核心模块未加载")
    result = to.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),
):
    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)
    ed = _get_cached("edm_designer")
    if not ed:
        raise HTTPException(503, "服务不可用：核心模块未加载")
    result = ed.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)
    ms = _get_cached("machining_simulator")
    if not ms:
        raise HTTPException(503, "服务不可用：核心模块未加载")
    result = ms.simulate_machining(operations, stock_bbox, resolution)
    return {"status": "success", "data": result}


@router.post("/export-mold")
async def export_mold_results(
    request: Request,
    current_user: User = Depends(get_current_active_user),
    db_session: AsyncSession = Depends(get_db_session),
):
    body = await request.json()
    task_id = body.get("task_id")
    scheme_id = body.get("scheme_id")
    formats = body.get("formats", ["step", "stl"])
    components = body.get("components", ["cavity", "core"])

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

    await _ensure_task_access(db_session, task_id, current_user.id)
    task_data = await TaskQueryService.get_task_view(db_session, task_id)
    if not task_data:
        raise HTTPException(404, "任务不存在")

    resolved_scheme_id = scheme_id or task_data.get("best_scheme_id") or "default"
    persisted_files = _select_persisted_files(
        task_id=task_id,
        task_data=task_data,
        scheme_id=resolved_scheme_id,
        formats=formats,
        components=components,
    )
    if persisted_files:
        return {
            "status": "success",
            "data": {
                "base_filename": Path(task_data.get("filename", f"mold_{task_id}")).stem,
                "task_id": task_id,
                "scheme_id": resolved_scheme_id,
                "files": persisted_files,
                "errors": [],
                "total_files": len(persisted_files),
                "total_errors": 0,
                "source": "persisted",
            },
        }

    cavity_shapes = processing_service.get_export_shapes(
        task_id,
        resolved_scheme_id,
    )
    filename = task_data.get("filename", f"mold_{task_id}")

    if not cavity_shapes:
        raise HTTPException(
            409,
            "导出缓存已失效或任务尚未完成，请重新分析后再导出以保证方案一致性",
        )

    base_filename = Path(filename).stem
    result = cad_exporter.export_mold_results(
        cavity_data=cavity_shapes,
        base_filename=base_filename,
        formats=formats,
        components=components,
        task_id=task_id,
        scheme_id=resolved_scheme_id,
    )
    result["files"] = _augment_export_files(task_id, result.get("files", []))
    result["source"] = "generated"

    merged_artifacts = _merge_export_artifacts(_get_export_artifacts(task_data), result)
    await storage_service.update_task_parameters(
        db_session,
        task_id,
        {"export_artifacts": merged_artifacts},
    )
    await redis_task_manager.update_task(task_id, {"export_artifacts": merged_artifacts})

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


@router.get("/export-download/{filepath:path}")
async def download_export_file(
    filepath: str,
    task_id: str,
    current_user: User = Depends(get_current_active_user),
    db_session: AsyncSession = Depends(get_db_session),
):
    from fastapi.responses import FileResponse

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

    await _ensure_task_access(db_session, task_id, current_user.id)
    task_data = await TaskQueryService.get_task_view(db_session, task_id)
    if not task_data:
        raise HTTPException(404, "任务不存在")

    allowed_paths = set()
    artifacts = _get_export_artifacts(task_data)
    for scheme in (artifacts.get("schemes") or {}).values():
        for file in scheme.get("files", []):
            relative_path = str(file.get("relative_path") or "").replace("\\", "/").strip("/")
            if relative_path:
                allowed_paths.add(relative_path)

    normalized_path = str(filepath or "").replace("\\", "/").strip("/")
    if normalized_path not in allowed_paths:
        raise HTTPException(403, "该文件不在任务允许下载清单中")

    full_path = os.path.join(cad_exporter.output_dir, normalized_path.replace("/", os.sep))
    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),
):
    result = cad_exporter.get_export_recommendations(target)
    return {"status": "success", "data": result}