import spaces
import os
import gradio as gr
from transformers import pipeline
from fastapi import HTTPException
from pydantic import BaseModel, ConfigDict
from typing import List

# ── Model ──────────────────────────────────────────────────────────────────────
ner = pipeline(
    "ner",
    model="Ramosh99/mixcode_pii_detector",
    aggregation_strategy="average",
    token=os.getenv("HF_TOKEN"),
)


# ── Core inference (ZeroGPU — GPU used only during inference call) ──────────────
@spaces.GPU
def detect_pii(text: str):
    """Run NER and return structured PII results (used by REST API)."""
    if not text.strip():
        return {"text": text, "entities": []}

    raw = ner(text)
    entities = []
    for item in raw:
        label = item.get("entity_group", "")
        if label.startswith(("B-", "I-")):
            label = label[2:]
        entities.append({
            "entity": label,
            "start":  item["start"],
            "end":    item["end"],
            "score":  round(float(item["score"]), 4),
        })
    return {"text": text, "entities": entities}


def detect_pii_ui(text: str):
    """Return list-of-tuples for Gradio HighlightedText (used by UI)."""
    res = detect_pii(text)
    entities = sorted(res["entities"], key=lambda x: x["start"])

    if not entities:
        return [(text, None)]

    result, last = [], 0
    for ent in entities:
        if ent["start"] > last:
            result.append((text[last:ent["start"]], None))
        result.append((text[ent["start"]:ent["end"]], ent["entity"]))
        last = ent["end"]
    if last < len(text):
        result.append((text[last:], None))
    return result


# ── Gradio Interface ───────────────────────────────────────────────────────────
demo = gr.Interface(
    fn=detect_pii_ui,
    inputs=gr.Textbox(
        label="Input Text",
        placeholder="Enter Sinhala-English codemixed text...",
        lines=4,
    ),
    outputs=gr.HighlightedText(label="Detected PII Entities"),
    title="Sinhala-English Codemix PII Detector",
    description=(
        "Detects PII entities (PERSON, ORG, PHONE, EMAIL, NIC, etc.) in "
        "Sinhala-English code-mixed text.\n\n"
        "**REST API:** `POST /api/detect` · `GET /api/health` · Docs: `/docs`"
    ),
    examples=[
        ["Mage name Kamal, man wada karanne MAS Holdings eke, phone 071-1234567"],
        ["මගේ email address kamal.perera@gmail.com, NIC number 199912345678"],
        ["I live in Colombo and my phone number is 0771234567"],
    ],
    flagging_mode="never",
)

# ── Launch Gradio — properly registers @spaces.GPU with ZeroGPU ───────────────
gr_app, _local_url, _share_url = demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    prevent_thread_lock=True,
    show_api=False,
)

# ── REST API routes (added to Gradio's internal FastAPI app) ───────────────────

class DetectRequest(BaseModel):
    model_config = ConfigDict(
        json_schema_extra={"example": {"text": "Mage name Kamal, phone 071-1234567"}}
    )
    text: str


class Entity(BaseModel):
    entity: str
    start:  int
    end:    int
    score:  float


class DetectResponse(BaseModel):
    text:          str
    entities:      List[Entity]
    redacted_text: str


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


@gr_app.post("/api/detect", response_model=DetectResponse)
async def api_detect_pii(request: DetectRequest):
    if not request.text.strip():
        raise HTTPException(status_code=422, detail="Input text cannot be empty.")

    result   = detect_pii(request.text)
    text     = result["text"]
    entities = sorted(result["entities"], key=lambda e: e["start"], reverse=True)

    redacted = text
    for ent in entities:
        redacted = redacted[: ent["start"]] + f"[{ent['entity']}]" + redacted[ent["end"]:]

    return {
        "text":          result["text"],
        "entities":      result["entities"],
        "redacted_text": redacted,
    }


# ── Keep process alive ─────────────────────────────────────────────────────────
demo.block_thread()