from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, WebSocket, Response from fastapi.responses import FileResponse from sqlalchemy.orm import Session from typing import List import os from ...core.config import settings from ...core.database import get_db from ...models.file import File as FileModel, FileStatus from ...services.file_service import process_file, delete_file from ...schemas.file import FileResponse as FileResponseSchema, FileList import asyncio from fastapi import WebSocketDisconnect import uuid router = APIRouter() @router.post("/upload", response_model=FileResponseSchema) async def upload_file( file: UploadFile = File(...), db: Session = Depends(get_db) ): if not file.filename: raise HTTPException(status_code=400, detail="No file provided") if not any(file.filename.lower().endswith(ext) for ext in settings.ALLOWED_EXTENSIONS): raise HTTPException( status_code=400, detail=f"File type not allowed. Allowed types: {', '.join(settings.ALLOWED_EXTENSIONS)}" ) # Generate unique file ID file_id = str(uuid.uuid4()) file_extension = os.path.splitext(file.filename)[1] unique_filename = f"{file_id}{file_extension}" # Save file with unique name file_path = settings.UPLOAD_FOLDER / unique_filename with open(file_path, "wb") as buffer: content = await file.read() buffer.write(content) # Create database entry db_file = FileModel( id=file_id, filename=file.filename, original_path=str(file_path), status=FileStatus.NOT_STARTED ) db.add(db_file) db.commit() db.refresh(db_file) # Start processing process_file.delay(str(db_file.id)) return db_file @router.get("/files", response_model=List[FileResponseSchema]) def list_files( skip: int = 0, limit: int = 100, db: Session = Depends(get_db) ): files = db.query(FileModel).offset(skip).limit(limit).all() return files @router.get("/files/{file_id}", response_model=FileResponseSchema) def get_file( file_id: str, db: Session = Depends(get_db) ): file = db.query(FileModel).filter(FileModel.id == file_id).first() if not file: raise HTTPException(status_code=404, detail="File not found") return file @router.get("/files/{file_id}/download") async def download_file( file_id: str, db: Session = Depends(get_db) ): file = db.query(FileModel).filter(FileModel.id == file_id).first() if not file: raise HTTPException(status_code=404, detail="File not found") if file.status != FileStatus.SUCCESS: raise HTTPException(status_code=400, detail="File is not ready for download") if not os.path.exists(file.processed_path): raise HTTPException(status_code=404, detail="Processed file not found") return FileResponse( path=file.processed_path, filename=file.filename, media_type="application/octet-stream" ) @router.websocket("/ws/status/{file_id}") async def websocket_endpoint(websocket: WebSocket, file_id: str, db: Session = Depends(get_db)): await websocket.accept() try: while True: file = db.query(FileModel).filter(FileModel.id == file_id).first() if not file: await websocket.send_json({"error": "File not found"}) break await websocket.send_json({ "status": file.status, "error": file.error_message }) if file.status in [FileStatus.SUCCESS, FileStatus.FAILED]: break await asyncio.sleep(1) except WebSocketDisconnect: pass @router.delete("/files/{file_id}") async def delete_file_endpoint( file_id: str, db: Session = Depends(get_db) ): """ Delete a file and its associated records. This will remove: 1. The database record 2. The original uploaded file 3. The processed markdown file (if it exists) """ try: delete_file(file_id) return {"message": "File deleted successfully"} except HTTPException as e: raise e except Exception as e: raise HTTPException(status_code=500, detail=str(e))