legal-doc-masker/backend/app/api/endpoints/files.py

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)
):
    print(f"=== DOWNLOAD REQUEST ===")
    print(f"File ID: {file_id}")
    
    file = db.query(FileModel).filter(FileModel.id == file_id).first()
    if not file:
        print(f"❌ File not found for ID: {file_id}")
        raise HTTPException(status_code=404, detail="File not found")
    
    print(f"✅ File found: {file.filename}")
    print(f"File status: {file.status}")
    print(f"Original path: {file.original_path}")
    print(f"Processed path: {file.processed_path}")
    
    if file.status != FileStatus.SUCCESS:
        print(f"❌ File not ready for download. Status: {file.status}")
        raise HTTPException(status_code=400, detail="File is not ready for download")
    
    if not os.path.exists(file.processed_path):
        print(f"❌ Processed file not found at: {file.processed_path}")
        raise HTTPException(status_code=404, detail="Processed file not found")
    
    print(f"✅ Processed file exists at: {file.processed_path}")
    
    # Get the original filename without extension and add .md extension
    original_filename = file.filename
    filename_without_ext = os.path.splitext(original_filename)[0]
    download_filename = f"{filename_without_ext}.md"
    
    print(f"Original filename: {original_filename}")
    print(f"Filename without extension: {filename_without_ext}")
    print(f"Download filename: {download_filename}")
    
    
    response = FileResponse(
        path=file.processed_path,
        filename=download_filename,
        media_type="text/markdown"
    )
    
    print(f"Response headers: {dict(response.headers)}")
    print(f"=== END DOWNLOAD REQUEST ===")
    
    return response

@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))