# app.py

import os
import io
import uvicorn
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
from pydantic import BaseModel
from typing import List

# For simplicity, we directly use the fast functions.
# The `image_processor.py` file now contains the optimized versions.
from image_processor import enhance_image_fast, extract_text_from_image_fast, process_pdf_in_parallel

app = FastAPI(
    title="High-Speed OCR API",
    description="An API to extract text from images and PDFs, optimized for speed.",
    version="4.0.0-hf"
)

# CORS Middleware to allow requests from any origin
origins = ["*"]
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Pydantic Models for structured responses
class ImageOCRResponse(BaseModel):
    filename: str
    text: str

class PageResult(BaseModel):
    page_number: int
    text: str

class PDFOCRResponse(BaseModel):
    filename: str
    total_pages: int
    results: List[PageResult]

# API Endpoints
@app.get("/", tags=["General"])
def read_root():
    return {"message": "Welcome to the High-Speed OCR API. See /docs for documentation."}

@app.post("/ocr-image", response_model=ImageOCRResponse, tags=["OCR"])
async def ocr_image_endpoint(file: UploadFile = File(...)):
    """Accepts an image, enhances it, and returns the extracted text."""
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="File must be an image.")
    try:
        contents = await file.read()
        image = Image.open(io.BytesIO(contents))
        enhanced_image = enhance_image_fast(image)
        text = extract_text_from_image_fast(enhanced_image)
        return {"filename": file.filename, "text": text}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing image OCR: {e}")

@app.post("/ocr-pdf", response_model=PDFOCRResponse, tags=["OCR"])
async def ocr_pdf_endpoint(file: UploadFile = File(...)):
    """Accepts a PDF, extracts text from each page in parallel, and returns structured results."""
    if file.content_type != "application/pdf":
        raise HTTPException(status_code=400, detail="File must be a PDF.")
    try:
        contents = await file.read()
        results = process_pdf_in_parallel(contents)
        return {
            "filename": file.filename,
            "total_pages": len(results),
            "results": results
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing PDF: {e}")

# This block allows running the app locally for testing
if __name__ == "__main__":
    # Hugging Face Spaces expects the app to run on port 7860
    port = int(os.environ.get("PORT", 7860)) 
    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=True)