import os
import requests
from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import JSONResponse
from PIL import Image
from io import BytesIO
import base64

app = FastAPI()

HF_API_URL = "https://api-inference.huggingface.co/models/google/medgemma-27b-it"
HF_TOKEN = os.getenv("HF_TOKEN")

headers = {
    "Authorization": f"Bearer {HF_TOKEN}",
    "Content-Type": "application/json"
}

@app.post("/analyze/")
async def analyze_image(prompt: str = Form(...), file: UploadFile = File(...)):
    try:
        image_data = await file.read()
        image = Image.open(BytesIO(image_data)).convert("RGB")

        buffered = BytesIO()
        image.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode()

        payload = {
            "inputs": [
                {
                    "role": "user",
                    "content": [
                        {"type": "image", "image": img_str},
                        {"type": "text", "text": prompt}
                    ]
                }
            ]
        }

        response = requests.post(HF_API_URL, headers=headers, json=payload)

        if response.status_code != 200:
            return JSONResponse(status_code=response.status_code, content=response.json())

        return {"result": response.json()[0]['generated_text']}

    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})