Newer
Older
CNCTools / ReferenceSurfaceGenerator / backend / app / main.py
import asyncio
import os
import uuid
import importlib
from typing import List
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, status
from fastapi.responses import FileResponse
from starlette.websockets import WebSocketState
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from .job_manager import load_job_metadata, load_all_job_metadata, save_job_metadata, _get_job_metadata_path
from .dxf_parser import parse_dxf_for_viewing
from .models import Job

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- Feature Loading ---
def load_features():
    features_dir = os.path.join(os.path.dirname(__file__), "features")
    for feature_name in os.listdir(features_dir):
        if os.path.isdir(os.path.join(features_dir, feature_name)):
            try:
                module = importlib.import_module(f".features.{feature_name}.main", package=__package__)
                if hasattr(module, "router"):
                    app.include_router(module.router, prefix=f"/api/features/{feature_name}", tags=[feature_name])
                    print(f"Successfully loaded feature: {feature_name}")
            except ImportError as e:
                print(f"Failed to load feature {feature_name}: {e}")

load_features()

# --- WebSocket Job Tracking ---
async def track_job_progress(websocket: WebSocket, job_id: uuid.UUID):
    initial_job = load_job_metadata(job_id)
    if not initial_job:
        await websocket.send_json({"status": "error", "message": "Job not found."})
        return

    last_update_content = initial_job.model_dump(mode='json')
    await websocket.send_json(last_update_content)

    while websocket.client_state == WebSocketState.CONNECTED:
        job = load_job_metadata(job_id)
        if not job:
            await websocket.send_json({"status": "error", "message": "Job disappeared or was deleted."})
            break
        
        update_content = job.model_dump(mode='json')
        if update_content != last_update_content:
            await websocket.send_json(update_content)
            last_update_content = update_content

        if job.status in ["COMPLETE", "FAILED"]:
            break
        
        await asyncio.sleep(0.5)

@app.websocket("/ws/{job_id}")
async def websocket_endpoint(websocket: WebSocket, job_id: uuid.UUID):
    await websocket.accept()
    try:
        await track_job_progress(websocket, job_id)
    except WebSocketDisconnect:
        print(f"Client disconnected from job {job_id}")
    finally:
        print(f"WebSocket connection handler finished for job {job_id}")

# --- Generic Job and File Management API Endpoints ---
@app.get("/api/jobs", response_model=List[Job])
async def get_all_jobs():
    return load_all_job_metadata()

@app.get("/api/jobs/{job_id}", response_model=Job)
async def get_job_status(job_id: uuid.UUID):
    job = load_job_metadata(job_id)
    if not job:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Job not found")
    return job

@app.delete("/api/jobs/{job_id}", status_code=status.HTTP_204_NO_CONTENT)
async def delete_job(job_id: uuid.UUID):
    job = load_job_metadata(job_id)
    if not job:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Job not found")

    if job.input_path and os.path.exists(job.input_path):
        os.remove(job.input_path)
    if job.output_path and os.path.exists(job.output_path):
        os.remove(job.output_path)
    
    os.remove(_get_job_metadata_path(job_id))
@app.get("/api/jobs/{job_id}/view")
async def get_job_output_for_viewing(job_id: uuid.UUID):
    """
    Retrieves the geometric data from a job's output DXF file in a web-friendly JSON format.
    """
    job = load_job_metadata(job_id)
    if not job:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Job not found")

    if not job.output_path or not os.path.exists(job.output_path):
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Output file not found for this job")

    try:
        data = parse_dxf_for_viewing(job.output_path)
        return data
    except IOError:
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Could not read the DXF file.")
    except Exception as e:
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"An unexpected error occurred while parsing the DXF file: {e}")


@app.get("/api/download/{job_id}")
async def download_file(job_id: uuid.UUID):
    job = load_job_metadata(job_id)
    if not job or not job.output_path or not os.path.exists(job.output_path):
        raise HTTPException(status_code=404, detail="File not found")
    
    filename = os.path.basename(job.output_path)
    return FileResponse(job.output_path, media_type='application/octet-stream', filename=filename)

# --- Static Files Hosting ---
if os.path.isdir("/app/static"):
    app.mount("/", StaticFiles(directory="/app/static", html=True), name="static")

@app.get("/api/features")
async def get_features():
    """
    Returns a list of available features.
    """
    features_dir = os.path.join(os.path.dirname(__file__), "features")
    return [
        f for f in os.listdir(features_dir) 
        if os.path.isdir(os.path.join(features_dir, f)) and not f.startswith("__")
    ]