import gradio as gr
from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
import google.generativeai as genai
import base64
import json
import uvicorn
from typing import Dict
from PIL import Image
import io

GOOGLE_API_KEY = 'AIzaSyDcYyq3w21iwipYn17wCAQo3AYWhUIGDSI'
genai.configure(api_key=GOOGLE_API_KEY)
model = genai.GenerativeModel(model_name="gemini-1.5-pro")

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

def process_image(image) -> Dict:
    if isinstance(image, str):
        image = Image.open(image)
    
    # Convert PIL Image to bytes
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format='JPEG')
    img_byte_arr = img_byte_arr.getvalue()
    
    encoded_image = base64.b64encode(img_byte_arr).decode("utf-8")
    
    prompt = """
    Analyze the business card and extract the following information in JSON format:
    {
      "name": "Full Name",
      "mobile_number": "Mobile Number",
      "company_name": "Company Name",
      "email": "Email Address",
      "job_title": "Job Title",
      "address": "Address",
      "website": "Website URL"
    }
    Ensure that the response is strictly in JSON format.
    """
    
    response = model.generate_content([
        {"mime_type": "image/jpeg", "data": encoded_image},
        prompt,
    ])
    
    raw = response.text.strip()
    if raw.startswith("```"):
        lines = raw.splitlines()
        if lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].startswith("```"):
            lines = lines[:-1]
        raw = "\n".join(lines).strip()
    
    return json.loads(raw)

@app.post("/extract")
async def extract_info(file: UploadFile = File(...)):
    contents = await file.read()
    img = Image.open(io.BytesIO(contents))
    result = process_image(img)
    return result

def gradio_process(image):
    if image is None:
        return json.dumps({"error": "No image provided"}, indent=2)
    result = process_image(image)
    return json.dumps(result, indent=2)

interface = gr.Interface(
    fn=gradio_process,
    inputs=gr.Image(type="pil", label="Upload Business Card"),
    outputs=gr.JSON(label="Extracted Information"),
    title="Business Card Information Extractor",
    description="Upload a business card image to extract contact information"
)

app = gr.mount_gradio_app(app, interface, path="/")

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