from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
from gliner import GLiNER
import os
import time

app = FastAPI(title="SIGAP AI ML Service (HF Spaces)")

MODEL_HUB = "urchade/gliner_multi-v2.1"

print(f"Downloading model {MODEL_HUB} from HuggingFace...", flush=True)

try:
    model = GLiNER.from_pretrained(MODEL_HUB)
    print("Model GLiNER loaded successfully!", flush=True)
except Exception as e:
    print(f"Failed loading GLiNER model: {e}", flush=True)
    raise e

LABELS = ["nama", "usia", "kondisi medis", "keterbatasan mobilitas", "asal lokasi", "anggota keluarga", "ketiadaan obat"]

class ExtractRequest(BaseModel):
    teks: str

class Entity(BaseModel):
    label: str
    text: str
    confidence: float

@app.post("/extract")
async def extract_entities(req: ExtractRequest):
    start_time = time.time()
    try:
        predictions = model.predict_entities(req.teks, LABELS, threshold=0.4)
        entities = [{"label": p["label"], "text": p["text"], "confidence": float(p["score"])} for p in predictions]
        latency = (time.time() - start_time) * 1000
        
        return {"entities": entities, "latency_ms": latency}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
async def root():
    return {"status": "Model GLiNER Ready!"}