import os
import io
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
import google.generativeai as genai
import uvicorn

app = FastAPI()

# CORS Middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

# Configure Gemini API
genai.configure(api_key=os.getenv("GEMINI_API_KEY"))

def extract_mixed_text(image):
    """Extract text from the uploaded image."""
    try:
        model = genai.GenerativeModel('gemini-2.0-flash')
        response = model.generate_content(
            [
                "CRITICAL INSTRUCTIONS FOR TEXT EXTRACTION: "
                "1. Extract ALL text from this image. "
                "2. Include Malayalam, English letters, and Numbers. "
                "3. Preserve formatting and order.",
                image
            ]
        )
        extracted_text = response.text

        # Filter text to remove unwanted symbols
        filtered_text = ''.join(
            char for char in extracted_text 
            if ('\u0D00' <= char <= '\u0D7F') or 
               char.isalnum() or 
               char.isspace() or 
               char in '.,/-:;()[]%₹$\n'
        )
        return filtered_text.strip() if filtered_text.strip() else "No valid text detected."
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error extracting text: {str(e)}")

@app.post("/extract_text")
async def extract_text(image: UploadFile = File(...)):
    """API Endpoint for text extraction."""
    try:
        # Read the image file
        image_contents = await image.read()
        pil_image = Image.open(io.BytesIO(image_contents))
        
        # Extract text
        extracted_text = extract_mixed_text(pil_image)
        
        return {"extracted_text": extracted_text}
    
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

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