main.py

import os
import time
import tempfile
import traceback
from pathlib import Path
from typing import List, Union, Optional, Dict, Any, Literal

from fastapi import FastAPI, File, UploadFile, HTTPException, status
from pydantic import BaseModel, Field
from PIL import Image 

# --- IMPORTS ---
from mdr_pdf_parser import (
    MagicPDFProcessor,
    MDRStructuredBlock,
    MDRTextBlock,
    MDRTableBlock,
    MDRFormulaBlock,
    MDRFigureBlock,
    MDRTextKind,
    MDRTableFormat,
    MDRRectangle,
    MDRTextSpan,
    MDRExtractedTableFormat # For configuration
)

# --- Configuration ---
# Read from environment variables, falling back to defaults
MODEL_DIR = os.environ.get("MDR_MODEL_DIR", "/models") # Default path inside container
DEVICE = os.environ.get("MDR_DEVICE", "cuda") # Default to cuda, processor will fallback if needed
TABLE_FORMAT_STR = os.environ.get("MDR_TABLE_FORMAT", "MARKDOWN") # Default table format

# Convert table format string to Enum
try:
    TABLE_FORMAT = MDRExtractedTableFormat[TABLE_FORMAT_STR.upper()]
except KeyError:
    print(f"Warning: Invalid MDR_TABLE_FORMAT '{TABLE_FORMAT_STR}'. Defaulting to DISABLE.")
    TABLE_FORMAT = MDRExtractedTableFormat.DISABLE

# --- Initialize Processor ---
# This happens once when the service starts
print("Initializing MagicPDFProcessor...")
print(f"Model Directory: {MODEL_DIR}")
print(f"Target Device: {DEVICE}")
print(f"Table Format: {TABLE_FORMAT.name}")
start_time = time.time()
try:
    mdr_processor = MagicPDFProcessor(
        device=DEVICE,
        model_dir_path=MODEL_DIR,
        extract_table_format=TABLE_FORMAT,
        # Set debug_dir_path=None for production service
        debug_dir_path=None
    )
    print(f"MagicPDFProcessor initialized successfully ({time.time() - start_time:.2f}s)")
except Exception as e:
    print(f"FATAL ERROR: Failed to initialize MagicPDFProcessor during startup: {e}")
    print("Service cannot start.")
    traceback.print_exc()
    # Optionally exit or raise to prevent FastAPI from starting incorrectly
    mdr_processor = None # Ensure processor is None if init fails

# --- API Models (Pydantic) ---
# Define models for API input/output for validation and documentation

class MDRPointModel(BaseModel):
    x: float
    y: float

class MDRRectangleModel(BaseModel):
    lt: MDRPointModel
    rt: MDRPointModel
    lb: MDRPointModel
    rb: MDRPointModel

    @classmethod
    def from_mdr_rectangle(cls, rect: MDRRectangle):
        return cls(
            lt=MDRPointModel(x=rect.lt[0], y=rect.lt[1]),
            rt=MDRPointModel(x=rect.rt[0], y=rect.rt[1]),
            lb=MDRPointModel(x=rect.lb[0], y=rect.lb[1]),
            rb=MDRPointModel(x=rect.rb[0], y=rect.rb[1]),
        )

class MDRTextSpanModel(BaseModel):
    content: str
    rank: float
    rect: MDRRectangleModel

    @classmethod
    def from_mdr_text_span(cls, span: MDRTextSpan):
        return cls(
            content=span.content,
            rank=span.rank,
            rect=MDRRectangleModel.from_mdr_rectangle(span.rect)
        )

class MDRBasicBlockModel(BaseModel):
    block_type: str # To distinguish block types in the union
    rect: MDRRectangleModel
    texts: List[MDRTextSpanModel] = Field(default_factory=list) # Captions/footnotes
    font_size: float

class MDRTextBlockModel(MDRBasicBlockModel):
    block_type: Literal["TextBlock"] = "TextBlock"
    kind: str # Use string representation of enum
    has_paragraph_indentation: bool
    last_line_touch_end: bool
    # Override texts field specifically for TextBlock
    texts: List[MDRTextSpanModel] # Text content itself

    @classmethod
    def from_mdr_text_block(cls, block: MDRTextBlock):
        return cls(
            rect=MDRRectangleModel.from_mdr_rectangle(block.rect),
            texts=[MDRTextSpanModel.from_mdr_text_span(span) for span in block.texts],
            font_size=block.font_size,
            kind=block.kind.name, # Convert enum to string name
            has_paragraph_indentation=block.has_paragraph_indentation,
            last_line_touch_end=block.last_line_touch_end
        )

class MDRTableBlockModel(MDRBasicBlockModel):
    block_type: Literal["TableBlock"] = "TableBlock"
    content: str
    format: str # Use string representation of enum
    # Omit 'image' field from API response

    @classmethod
    def from_mdr_table_block(cls, block: MDRTableBlock):
        return cls(
            rect=MDRRectangleModel.from_mdr_rectangle(block.rect),
            texts=[MDRTextSpanModel.from_mdr_text_span(span) for span in block.texts], # Captions
            font_size=block.font_size,
            content=block.content,
            format=block.format.name # Convert enum to string name
        )

class MDRFormulaBlockModel(MDRBasicBlockModel):
    block_type: Literal["FormulaBlock"] = "FormulaBlock"
    content: Optional[str] = None
    # Omit 'image' field from API response

    @classmethod
    def from_mdr_formula_block(cls, block: MDRFormulaBlock):
        return cls(
            rect=MDRRectangleModel.from_mdr_rectangle(block.rect),
            texts=[MDRTextSpanModel.from_mdr_text_span(span) for span in block.texts], # Captions
            font_size=block.font_size,
            content=block.content
        )

class MDRFigureBlockModel(MDRBasicBlockModel):
    block_type: Literal["FigureBlock"] = "FigureBlock"
    # Omit 'image' field from API response

    @classmethod
    def from_mdr_figure_block(cls, block: MDRFigureBlock):
        return cls(
            rect=MDRRectangleModel.from_mdr_rectangle(block.rect),
            texts=[MDRTextSpanModel.from_mdr_text_span(span) for span in block.texts], # Captions
            font_size=block.font_size
        )

# Union type for the response model
MDRStructuredBlockModelAPI = Union[MDRTextBlockModel, MDRTableBlockModel, MDRFormulaBlockModel, MDRFigureBlockModel] # Renamed API Union type

# --- FastAPI App ---
app = FastAPI(
    title="MagicDataReadiness PDF Processor",
    description="API service to extract structured content from PDF files.",
    version="1.0.0"
)

@app.on_event("startup")
async def startup_event():
    if mdr_processor is None:
        # This prevents the app from starting if initialization failed
        raise RuntimeError("MagicPDFProcessor failed to initialize. Service cannot start.")
    print("MagicDataReadiness Service is ready.")

@app.get("/health")
async def health_check():
    """Simple health check endpoint."""
    if mdr_processor is None:
        raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Processor not initialized")
    return {"status": "ok", "message": "MagicPDFProcessor is running."}

@app.post("/process-pdf/",
          response_model=List[MDRStructuredBlockModelAPI], # Use the Union type
          summary="Process a PDF file",
          description="Upload a PDF file to extract structured blocks (text, tables, figures, formulas).")
async def process_pdf_endpoint(file: UploadFile = File(..., description="The PDF file to process.")):
    """
    Handles PDF file upload, processing, and returns extracted blocks.
    """
    if mdr_processor is None:
        raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Processor not initialized")

    if not file.filename.lower().endswith(".pdf"):
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid file type. Please upload a PDF.")

    # Save uploaded file temporarily
    temp_pdf_path = "" # Initialize path
    try:
        # Create a temporary directory if it doesn't exist
        temp_dir = Path("./temp_uploads")
        temp_dir.mkdir(exist_ok=True)

        # Use a temporary file within the directory
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf", dir=temp_dir) as temp_file:
            content = await file.read()
            temp_file.write(content)
            temp_pdf_path = temp_file.name
        print(f"Received file '{file.filename}', saved temporarily to '{temp_pdf_path}'")
    except Exception as e:
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to save uploaded file: {e}")

    extracted_blocks_api: List[MDRStructuredBlockModelAPI] = []
    start_process_time = time.time()

    try:
        print(f"Processing '{temp_pdf_path}'...")
        # Process the document using the temporary file path
        # Note: process_document returns a generator, collect all blocks
        all_blocks = list(mdr_processor.process_document(pdf_input=temp_pdf_path))
        print(f"Extracted {len(all_blocks)} raw blocks.")

        # Convert internal block types to API response models
        for block in all_blocks:
            if isinstance(block, MDRTextBlock):
                extracted_blocks_api.append(MDRTextBlockModel.from_mdr_text_block(block))
            elif isinstance(block, MDRTableBlock):
                extracted_blocks_api.append(MDRTableBlockModel.from_mdr_table_block(block))
            elif isinstance(block, MDRFormulaBlock):
                extracted_blocks_api.append(MDRFormulaBlockModel.from_mdr_formula_block(block))
            elif isinstance(block, MDRFigureBlock):
                extracted_blocks_api.append(MDRFigureBlockModel.from_mdr_figure_block(block))
            # Add more elif clauses if there are other block types

        process_time = time.time() - start_process_time
        print(f"Processing finished in {process_time:.2f}s. Returning {len(extracted_blocks_api)} blocks.")

    except Exception as e:
        print(f"ERROR during PDF processing: {e}")
        traceback.print_exc()
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"An error occurred during PDF processing: {e}")
    finally:
        # Clean up the temporary file
        if temp_pdf_path and os.path.exists(temp_pdf_path):
            try:
                os.remove(temp_pdf_path)
                print(f"Cleaned up temporary file: {temp_pdf_path}")
            except OSError as e:
                print(f"Warning: Could not remove temporary file {temp_pdf_path}: {e}")

    return extracted_blocks_api

# Optional: Add root endpoint for basic info/docs link
@app.get("/")
async def read_root():
    return {
        "message": "Welcome to the MagicDataReadiness PDF Processor API!",
        "docs_url": "/docs",
        "health_url": "/health"
    }