routers/upload.py

"""
Upload Router — Handles file uploads, extraction, and JSON response.

POST /api/upload
  - Accepts multipart/form-data with a single file
  - Validates file type, size, and user tier
  - Routes to PDF Parser or OCR Engine based on file analysis
  - Returns structured JSON with extracted fields
"""
import time
from pathlib import Path
from fastapi import APIRouter, UploadFile, File, Form, Header, HTTPException
from fastapi.responses import JSONResponse

from models.schemas import (
    ExtractionResponse,
    ExtractionMetadata,
    ExtractedField,
    FileType,
    ProcessingLane,
    DocumentType,
    ErrorResponse,
)
from services.file_router import detect_processing_lane, get_pdf_page_count
from services.pdf_parser import extract_from_pdf, extract_tables_as_fields
from services.ocr_engine import extract_from_image, extract_from_scanned_pdf
from services.json_mapper import map_text_to_fields
from services.summarizer import generate_summary
from services.tier_manager import check_upload_allowed, record_usage
from utils.helpers import (
    validate_file_type,
    generate_temp_filename,
    cleanup_temp_file,
    UPLOAD_DIR,
    MAX_FILE_SIZE_UNREGISTERED,
)

router = APIRouter(prefix="/api", tags=["upload"])


@router.post("/upload", response_model=ExtractionResponse)
async def upload_file(
    file: UploadFile = File(...),
    document_type: str = Form(default="form"),
    x_session_token: str = Header(default="anonymous"),
    x_user_registered: str = Header(default="false"),
):
    """
    Upload a file for extraction. Returns structured JSON with field-value pairs.

    Headers:
      - X-Session-Token: Session identifier from frontend localStorage
      - X-User-Registered: "true" if user is authenticated via Supabase
    """
    start_time = time.time()
    is_registered = x_user_registered.lower() == "true"
    doc_type_enum = DocumentType(document_type) if document_type in [e.value for e in DocumentType] else DocumentType.FORM

    # Validate file type
    file_type = validate_file_type(file.filename or "unknown", file.content_type)
    if not file_type:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported file type. Please upload PDF, JPG, or PNG files only.",
        )

    # Read file content to check size
    content = await file.read()
    file_size = len(content)

    # Check tier limits
    tier_check = check_upload_allowed(x_session_token, file_size, is_registered)
    if not tier_check.allowed:
        raise HTTPException(status_code=429, detail=tier_check.message)

    # Save to temp file
    temp_filename = generate_temp_filename(file.filename or "upload")
    temp_path = UPLOAD_DIR / temp_filename

    try:
        with open(temp_path, "wb") as f:
            f.write(content)

        # Detect processing lane
        lane = detect_processing_lane(temp_path, file_type)

        # Process the file
        fields: list[ExtractedField] = []
        page_count = 1
        raw_text = ""

        if lane == ProcessingLane.PDF_PARSER:
            pdf_result = extract_from_pdf(temp_path)
            page_count = pdf_result["page_count"]
            raw_text = pdf_result["raw_text"]

            if doc_type_enum == DocumentType.FORM:
                table_fields = extract_tables_as_fields(pdf_result.get("tables", []))
                fields = map_text_to_fields(raw_text=raw_text, tables=table_fields)
            else:
                fields = [ExtractedField(name="Extracted Text", value=raw_text, field_type="text", confidence=0.95)]

        elif lane == ProcessingLane.OCR_ENGINE:
            if file_type == "pdf":
                ocr_result = extract_from_scanned_pdf(temp_path)
                page_count = ocr_result["page_count"]
            else:
                ocr_result = extract_from_image(temp_path, preprocess=True)
                page_count = 1
                
            raw_text = ocr_result["raw_text"]

            if doc_type_enum == DocumentType.FORM:
                fields = map_text_to_fields(raw_text=raw_text, ocr_blocks=ocr_result.get("blocks", []))
            else:
                fields = [ExtractedField(name="Extracted Text", value=raw_text, field_type="text", confidence=0.95)]

        # Generate AI Summary
        summary = generate_summary(raw_text, is_registered=is_registered)

        # Record the usage
        record_usage(x_session_token)

        processing_time = int((time.time() - start_time) * 1000)

        # Calculate actual average confidence from extracted fields
        if fields:
            avg_confidence = sum(f.confidence for f in fields) / len(fields)
        else:
            avg_confidence = 0.0

        return ExtractionResponse(
            success=True,
            fields=fields,
            summary=summary,
            metadata=ExtractionMetadata(
                filename=file.filename or "unknown",
                file_type=FileType(file_type),
                processing_lane=lane,
                document_type=doc_type_enum,
                page_count=page_count,
                processing_time_ms=processing_time,
                confidence_score=round(avg_confidence * 100, 1),
            ),
            message=f"Extracted {len(fields)} fields from {page_count} page(s) in {processing_time}ms",
        )

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Processing error: {str(e)}",
        )
    finally:
        cleanup_temp_file(temp_path)


@router.get("/tier-status")
async def get_tier(
    x_session_token: str = Header(default="anonymous"),
    x_user_registered: str = Header(default="false"),
):
    """Get the current usage status for the session."""
    from services.tier_manager import get_tier_status

    is_registered = x_user_registered.lower() == "true"
    tier = get_tier_status(x_session_token, is_registered)
    return tier