"""
Sinhala-English Codemix PII Detector — HuggingFace Space API
Model: Ramosh99/mixcode_pii_detector (XLM-RoBERTa token-classification)
Also used as fallback for pure Sinhala inputs.
"""
import os
from transformers import pipeline
from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI(title="Codemix PII Detector", version="1.0.0")

# Load model at startup — stays in memory
ner = pipeline(
    "ner",
    model="Ramosh99/mixcode_pii_detector",
    aggregation_strategy="average",
    token=os.getenv("HF_TOKEN"),
)


class PredictRequest(BaseModel):
    text: str


@app.get("/health")
def health():
    return {"status": "ok", "model": "Ramosh99/mixcode_pii_detector"}


@app.post("/predict")
def predict(req: PredictRequest):
    raw = ner(req.text)
    entities = []
    for item in raw:
        label = item.get("entity_group", "")
        if label.startswith(("B-", "I-")):
            label = label[2:]
        entities.append({
            "span": item["word"].strip(),
            "start": item["start"],
            "end": item["end"],
            "entity_group": label,
            "confidence": round(float(item["score"]), 4),
        })
    return {"entities": entities}