backend/utils/validators.py

"""
Validation utilities for file upload pipeline.
"""
import os
from fastapi import UploadFile, HTTPException

STRUCTURED_EXTENSIONS = {"csv", "xlsx"}
UNSTRUCTURED_EXTENSIONS = {"pdf", "txt", "doc", "docx", "ppt", "pptx", "jpg", "jpeg", "png"}
ALLOWED_EXTENSIONS = STRUCTURED_EXTENSIONS | UNSTRUCTURED_EXTENSIONS

MIME_TYPE_MAP = {
    "csv": ["text/csv", "application/csv", "text/plain"],
    "xlsx": ["application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"],
    "pdf": ["application/pdf"],
    "txt": ["text/plain"],
    "doc": ["application/msword"],
    "docx": ["application/vnd.openxmlformats-officedocument.wordprocessingml.document"],
    "ppt": ["application/vnd.ms-powerpoint"],
    "pptx": ["application/vnd.openxmlformats-officedocument.presentationml.presentation"],
    "jpg": ["image/jpeg"],
    "jpeg": ["image/jpeg"],
    "png": ["image/png"],
}


def get_file_extension(filename: str) -> str:
    """Extract and normalize file extension."""
    if not filename or "." not in filename:
        return ""
    return filename.rsplit(".", 1)[-1].lower()


def validate_file_extension(filename: str) -> str:
    """
    Validate that file has an allowed extension.
    Returns the extension if valid, raises HTTPException otherwise.
    """
    ext = get_file_extension(filename)
    if not ext:
        raise HTTPException(
            status_code=400,
            detail="File must have a valid extension."
        )
    if ext not in ALLOWED_EXTENSIONS:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported file type: .{ext}. Allowed types: {', '.join(sorted(ALLOWED_EXTENSIONS))}"
        )
    return ext


def validate_file_size(file: UploadFile) -> None:
    """
    Validate that file is not empty.
    Raises HTTPException if file size is 0.
    """
    file.file.seek(0, os.SEEK_END)
    size = file.file.tell()
    file.file.seek(0)
    
    if size == 0:
        raise HTTPException(
            status_code=400,
            detail="File is empty. Please upload a valid file."
        )


def validate_mime_type(file: UploadFile, extension: str) -> None:
    """
    Light MIME type sanity check.
    Raises HTTPException if MIME type doesn't match extension.
    """
    content_type = file.content_type or ""
    expected_mimes = MIME_TYPE_MAP.get(extension, [])
    
    if not expected_mimes:
        return
    
    if content_type and content_type not in expected_mimes:
        if content_type != "application/octet-stream":
            raise HTTPException(
                status_code=400,
                detail=f"MIME type mismatch. Expected one of {expected_mimes} for .{extension}, got {content_type}"
            )


async def validate_upload_file(file: UploadFile) -> str:
    """
    Run all validations on uploaded file.
    Returns the validated file extension.
    """
    if not file or not file.filename:
        raise HTTPException(
            status_code=400,
            detail="No file provided or file has no name."
        )
    
    extension = validate_file_extension(file.filename)
    validate_file_size(file)
    validate_mime_type(file, extension)
    
    return extension