import asyncio
import base64
from concurrent.futures import ProcessPoolExecutor
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse

from compiler import compile_latex, CompilationResult

executor: ProcessPoolExecutor = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global executor
    executor = ProcessPoolExecutor(max_workers=4)
    yield
    executor.shutdown(wait=True)

app = FastAPI(lifespan=lifespan)

static_dir = Path(__file__).parent / "static"
app.mount("/static", StaticFiles(directory=static_dir), name="static")

@app.get("/")
async def index():
    return FileResponse(static_dir / "index.html")

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_json()
            if data.get("type") == "compile":
                latex_content = data.get("latex", "")
                await websocket.send_json({"type": "compiling"})
                
                loop = asyncio.get_event_loop()
                result: CompilationResult = await loop.run_in_executor(
                    executor, compile_latex, latex_content
                )
                
                if result.success:
                    pdf_base64 = base64.b64encode(result.pdf_data).decode("utf-8")
                    await websocket.send_json({"type": "pdf", "data": pdf_base64})
                else:
                    await websocket.send_json({"type": "error", "log": result.error_log})
    except WebSocketDisconnect:
        pass

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