kneron_model_converter/services/workers/nef/main.py

"""
NEF?????????
??????NEF?????????????????
"""

from fastapi import FastAPI, HTTPException, BackgroundTasks, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field
from typing import Optional, Dict, Any
import os
import logging
from datetime import datetime
from core import process_nef_core
from facade import WorkerFacade

# ????????
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="NEF Compilation Service",
    description="NEF?????????????",
    version="1.0.0"
)

# CORS???? - ??????Web?????????
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:4000"],  # ??????Web?????
    allow_credentials=True,
    allow_methods=["GET", "POST", "DELETE"],
    allow_headers=["*"],
)

# ????
UPLOAD_FOLDER = os.getenv("NEF_UPLOAD_FOLDER", "/tmp/nef_uploads")
OUTPUT_FOLDER = os.getenv("NEF_OUTPUT_FOLDER", "/tmp/nef_outputs")
API_KEY = os.getenv("NEF_API_KEY", "nef-secret-key")

# ??????????
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

worker = WorkerFacade(
    process_core=process_nef_core,
    output_folder=OUTPUT_FOLDER,
    output_prefix="compiled_",
    output_extension=".nef",
    step_name="nef",
    logger=logger,
)

# Pydantic????
class HealthResponse(BaseModel):
    service: str
    status: str
    timestamp: str
    active_tasks: int

class NEFProcessRequest(BaseModel):
    bie_file_id: str
    model_id: int = Field(..., ge=1, le=65535)
    version: str = Field(..., regex=r'^[0-9a-fA-F]{4}$')
    platform: str = Field(..., regex=r'^(520|720|530|630|730)$')

class TaskStatusResponse(BaseModel):
    task_id: str
    status: str
    progress: float
    message: str
    result: Optional[Dict[str, Any]] = None
    created_at: str
    updated_at: str

# ????????
def verify_api_key(x_api_key: str = None):
    """????PI Key"""
    if x_api_key != API_KEY:
        raise HTTPException(
            status_code=401,
            detail="Invalid API key"
        )
    return True

# ????????
@app.get("/health", response_model=HealthResponse)
async def health_check():
    """???????????"""
    active_tasks = worker.active_tasks_count()
    return HealthResponse(
        service="nef-compilation",
        status="healthy",
        timestamp=datetime.now().isoformat(),
        active_tasks=active_tasks
    )

# NEF???
@app.post("/api/nef/process")
async def process_nef(
    request: NEFProcessRequest,
    background_tasks: BackgroundTasks,
    api_key: str = Depends(verify_api_key)
):
    """???NEF????????"""
    try:
        # ????BIE???????????
        bie_file_path = os.path.join(UPLOAD_FOLDER, request.bie_file_id)
        if not os.path.exists(bie_file_path):
            raise HTTPException(status_code=404, detail="BIE file not found")

        task_id = worker.submit(
            parameters=request.dict(),
            input_paths={"bie_file_path": bie_file_path},
        )
        background_tasks.add_task(worker.run_task, task_id)

        logger.info(f"NEF compilation task submitted: {task_id}")

        return {
            "success": True,
            "task_id": task_id,
            "message": "NEF compilation task submitted successfully"
        }

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error processing NEF: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# ??????????
@app.get("/api/nef/tasks/{task_id}/status", response_model=TaskStatusResponse)
async def get_task_status(
    task_id: str,
    api_key: str = Depends(verify_api_key)
):
    """????EF???????"""
    task = worker.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")

    return TaskStatusResponse(
        task_id=task_id,
        status=task["status"],
        progress=task["progress"],
        message=task["message"],
        result=task["result"],
        created_at=task["created_at"].isoformat(),
        updated_at=task["updated_at"].isoformat()
    )

# ???????????
@app.get("/api/nef/tasks/{task_id}/result")
async def get_task_result(
    task_id: str,
    api_key: str = Depends(verify_api_key)
):
    """????EF???????"""
    task = worker.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")

    if task["status"] != "completed":
        raise HTTPException(status_code=400, detail="Task not completed yet")

    if "file_path" in task["result"]:
        return FileResponse(
            task["result"]["file_path"],
            media_type="application/octet-stream",
            filename=f"compiled_{task_id}.nef"
        )
    return task["result"]

# ???????
@app.get("/api/nef/tasks")
async def list_tasks(api_key: str = Depends(verify_api_key)):
    """???????EF????"""
    return {
        "tasks": [
            {
                "task_id": task["task_id"],
                "status": task["status"],
                "progress": task["progress"],
                "created_at": task["created_at"].isoformat(),
                "updated_at": task["updated_at"].isoformat()
            }
            for task in worker.list_tasks()
        ]
    }

# ????????
@app.delete("/api/nef/tasks/{task_id}")
async def cancel_task(
    task_id: str,
    api_key: str = Depends(verify_api_key)
):
    """????EF????"""
    try:
        worker.cancel_task(task_id)
    except KeyError:
        raise HTTPException(status_code=404, detail="Task not found")
    except ValueError:
        raise HTTPException(status_code=400, detail="Task cannot be cancelled")

    logger.info(f"NEF task cancelled: {task_id}")

    return {
        "success": True,
        "message": "NEF task cancelled successfully"
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5003)