from fastapi import FastAPI
from pydantic import BaseModel
import requests, base64, numpy as np, io, os
from PIL import Image

app = FastAPI(title="MedGemma ICD-10 Remote API")

HF_TOKEN = os.getenv("HF_TOKEN", "YOUR_HF_API_TOKEN")
MODEL_ID = "google/medgemma-4b-it"
API_URL = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
HEADERS = {"Authorization": f"Bearer {HF_TOKEN}", "Content-Type": "application/json"}

class TextInput(BaseModel):
    text: str

@app.get("/")
def root():
    return {"status": "running", "model": MODEL_ID, "api": API_URL}

@app.post("/predict")
def predict_icd10(input: TextInput):
    try:
        # dummy image
        dummy = Image.fromarray(np.zeros((224, 224, 3), dtype=np.uint8))
        buf = io.BytesIO()
        dummy.save(buf, format="PNG")
        image_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")

        prompt = f"<image>\nConvert this medical note into ICD-10 code only:\n{input.text}\nICD-10 code:"

        payload = {"inputs": {"image": image_b64, "prompt": prompt}}

        r = requests.post(API_URL, headers=HEADERS, json=payload, timeout=90)
        if r.status_code == 200:
            data = r.json()
            if isinstance(data, list) and "generated_text" in data[0]:
                return {"icd10_code": data[0]["generated_text"].strip()}
            return {"raw_response": data}
        else:
            return {"error": f"Hugging Face API error {r.status_code}", "details": r.text}
    except Exception as e:
        return {"error": str(e)}