import asyncio
import os
import uuid
from fastapi import FastAPI, File, UploadFile, WebSocket, WebSocketDisconnect, Request
from fastapi.responses import FileResponse
from starlette.websockets import WebSocketState

from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates

# Import the core processing logic
from .processing import create_layered_curves_dxf

app = FastAPI()

# Allow all origins for simplicity, can be locked down in production
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Create a directory for uploads and outputs
UPLOAD_DIR = "/app/data/uploads"
OUTPUT_DIR = "/app/data/outputs"
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)

# Mount the static files directory for the React app
@app.post("/upload/")
async def upload_mesh_file(file: UploadFile = File(...)):
    """
    Accepts a file upload and saves it to a temporary location.
    Returns a unique file ID to the client.
    """
    # Create a unique filename to avoid collisions
    file_id = str(uuid.uuid4())
    input_path = os.path.join(UPLOAD_DIR, f"{file_id}_{file.filename}")
    
    with open(input_path, "wb") as buffer:
        buffer.write(await file.read())
        
    return {"file_id": file_id, "filename": file.filename}


@app.websocket("/ws/{file_id}/{filename}")
async def websocket_endpoint(websocket: WebSocket, file_id: str, filename: str):
    """
    Handles the WebSocket connection for processing the file and sending progress.
    """
    await websocket.accept()
    
    input_path = os.path.join(UPLOAD_DIR, f"{file_id}_{filename}")
    output_filename = f"{file_id}_curves.dxf"
    output_path = os.path.join(OUTPUT_DIR, output_filename)

    try:
        # Call the actual processing function, which is now a generator
        async for progress in async_generator_wrapper(create_layered_curves_dxf(input_path, output_path)):
            if websocket.client_state != WebSocketState.DISCONNECTED:
                await websocket.send_json(progress)
            else:
                print("WebSocket disconnected, stopping process.")
                break
        
        if websocket.client_state != WebSocketState.DISCONNECTED:
            await websocket.send_json({
                "status": "complete",
                "progress": 100,
                "message": "Processing complete!",
                "download_url": f"/download/{output_filename}"
            })

    except Exception as e:
        error_message = f"An error occurred: {str(e)}"
        print(error_message)
        if websocket.client_state != WebSocketState.DISCONNECTED:
            await websocket.send_json({
                "status": "error",
                "message": error_message
            })
    finally:
        # Clean up the uploaded file
        if os.path.exists(input_path):
            os.remove(input_path)
        if websocket.client_state != WebSocketState.DISCONNECTED:
            await websocket.close()

async def async_generator_wrapper(generator):
    """
    A wrapper to run the synchronous generator in a separate thread
    to avoid blocking the asyncio event loop.
    """
    loop = asyncio.get_event_loop()
    for item in await loop.run_in_executor(None, list, generator):
        yield item


@app.get("/download/{filename}")
async def download_file(filename: str):
    """
    Serves the generated DXF file for download.
    """
    path = os.path.join(OUTPUT_DIR, filename)
    if os.path.exists(path):
        return FileResponse(path, media_type='application/vnd.dxf', filename=filename)
    return {"error": "File not found"}


# This MUST be the last route
app.mount("/", StaticFiles(directory="static", html = True), name="static")