from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel
from PIL import Image
import io
import os
import google.generativeai as genai
from pdf_generator import create_receipt


genai.configure(api_key=os.environ.get("GOOGLE_API_KEY"))

app = FastAPI()

class OCRResponse(BaseModel):
    extracted_text: str


@app.post("/ocr-image/", response_model=OCRResponse)
async def ocr_image(file: UploadFile = File(...)):
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="Uploaded file must be an image.")

    try:
        image_data = await file.read()
        image = Image.open(io.BytesIO(image_data))

        model = genai.GenerativeModel("gemini-2.0-flash")

        prompt = (
            "Extraire le texte de cette image et générer un reçu à propos les marchandises, "
            "sans aucune information supplémentaire."
        )

        contents = [prompt, image]
        response = model.generate_content(contents)

        extracted_text = response.text

        # If PDF generation is desired:
        create_receipt(extracted_text, file_name="result/receipt.pdf")
        return FileResponse("result/receipt.pdf", filename="receipt.pdf", media_type="application/pdf")

        

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"OCR failed: {str(e)}")