app.py

import sys
import os
import uuid
import tempfile
import pandas as pd
import torch
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from typing import List, Optional
from src.extractor import PassportExtractor
from src.formats import format_fly_dubai, format_iraqi, format_fly_baghdad

# Path fix for local imports
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

app = FastAPI()

# Initialize extractor
use_gpu = torch.cuda.is_available()
extractor = PassportExtractor(use_gpu=use_gpu)

# Temporary storage for download files
DOWNLOAD_DIR = os.path.join(tempfile.gettempdir(), "passport_downloads")
os.makedirs(DOWNLOAD_DIR, exist_ok=True)

@app.post("/api/process")
async def process_passports(
    files: List[UploadFile] = File(...),
    airline: str = Form(...),
    include_all: bool = Form(False)
):
    try:
        all_results = []
        problematic_files = []
        
        # Save uploaded files temporarily to process
        for file in files:
            file_name = file.filename
            content = await file.read()
            
            with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file_name)[1]) as tmp:
                tmp.write(content)
                tmp_path = tmp.name
            
            try:
                is_pdf = file_name.lower().endswith('.pdf')
                if is_pdf:
                    results = extractor.process_pdf(tmp_path, airline=airline.lower())
                    if not results:
                        results = [{"mrz_found": False}] # Placeholder for failure
                else:
                    result = extractor.get_data(tmp_path, airline=airline.lower())
                    results = [result] if result else []

                # Attach source file name
                mrz_found_in_file = False
                for res in results:
                    res['source_file'] = file_name
                    if res.get('mrz_found', False):
                        mrz_found_in_file = True
                
                all_results.extend(results)
                
                if not results or not mrz_found_in_file:
                    problematic_files.append({'file_name': file_name, 'reason': 'MRZ not found'})
                    
            except Exception as e:
                problematic_files.append({'file_name': file_name, 'reason': str(e)})
            finally:
                if os.path.exists(tmp_path):
                    os.remove(tmp_path)

        # Filter and format
        good_results = [r for r in all_results if not any(p['file_name'] == r.get('source_file') for p in problematic_files)]
        
        if airline == "Fly Dubai":
            df = format_fly_dubai(good_results)
        elif airline == "Iraqi":
            df = format_iraqi(good_results)
        elif airline == "Fly Baghdad":
            df = format_fly_baghdad(good_results)
        else:
            df = pd.DataFrame(good_results)

        # Add placeholders if requested
        if include_all and problematic_files:
            # Re-implementing simplified placeholder logic for API
            placeholders = []
            for p in problematic_files:
                placeholders.append({"Last Name": "•••", "First Name": "•••", "source_file": p['file_name']})
            df = pd.concat([df, pd.DataFrame(placeholders)], ignore_index=True)

        # Generate unique IDs for downloads
        session_id = str(uuid.uuid4())
        csv_name = f"passports_{session_id}.csv"
        xlsx_name = f"passports_{session_id}.xlsx"
        
        df.to_csv(os.path.join(DOWNLOAD_DIR, csv_name), index=False)
        df.to_excel(os.path.join(DOWNLOAD_DIR, xlsx_name), index=False)
        
        return {
            "data": df.fillna("").to_dict(orient="records"),
            "csv_id": csv_name,
            "xlsx_id": xlsx_name,
            "summary": {
                "total": len(files),
                "success": len(good_results),
                "failed": len(problematic_files)
            }
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/download/{file_name}")
async def download_file(file_name: str):
    file_path = os.path.join(DOWNLOAD_DIR, file_name)
    if os.path.exists(file_path):
        return FileResponse(file_path, filename=file_name)
    raise HTTPException(status_code=404, detail="File not found")

# Serve static files (HTML/JS)
app.mount("/", StaticFiles(directory="static", html=True), name="static")

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