"""
CLOAK PII Demo — ONNX INT8 CPU inference + Gradio UI + REST /detect API.
Downloads the model (MODEL_FILE) from HF Hub at startup.
Set HF_TOKEN secret in Space settings if the repo is private.
The /detect endpoint is what the CLOAK API (Next.js) calls via callHfSpace:
  POST /detect  Authorization: Bearer <API_SECRET>  {"text": "..."}
  -> {entities, nested, redacted_text, model, processing_time_ms}
"""

import html as _html
import json
import os
import time

# HF's Gradio-Space image sets GRADIO_SSR_MODE=true, which spawns a second
# (Node SSR) server on port 7861. Under the mount_gradio_app + uvicorn pattern
# it double-binds and crashes the Space ("[Errno 98] address already in use" on
# 7861). Force it off BEFORE gradio is imported.
os.environ["GRADIO_SSR_MODE"] = "false"

import gradio as gr
import numpy as np
import onnxruntime as ort
import pandas as pd
from fastapi import FastAPI, HTTPException, Request
from huggingface_hub import hf_hub_download, login
from pydantic import BaseModel, Field
from transformers import AutoTokenizer

# EntryNotFoundError is raised by hf_hub_download when a file isn't in the repo.
# Import location moved across huggingface_hub versions — try both.
try:
    from huggingface_hub.errors import EntryNotFoundError
except ImportError:  # older huggingface_hub
    from huggingface_hub.utils import EntryNotFoundError  # type: ignore

# ── Config ────────────────────────────────────────────────────────────────────

MODEL_REPO = os.getenv("MODEL_REPO", "Wardline/CLOAK_V3.10_onnx")
MODEL_FILE = os.getenv("MODEL_FILE", "model.fused.int8.onnx")  # fusion+int8
HF_TOKEN   = os.getenv("HF_TOKEN")
# Bearer token the CLOAK API must send. callHfSpace sends `Bearer ${HF_TOKEN}`,
# so default API_SECRET to HF_TOKEN — the same secret works for both unless set.
API_SECRET = os.getenv("API_SECRET", HF_TOKEN)

MAX_LEN   = 512
STRIDE    = 128
THRESHOLD = 0.0

# ── Thread count: read cgroup CPU quota, not sched_getaffinity ────────────────
# sched_getaffinity returns the host's full CPU set inside Docker containers.
# The real allocated vCPU count lives in the cgroup CFS quota.

def _container_cpu_count() -> int:
    # cgroup v2
    try:
        with open("/sys/fs/cgroup/cpu.max") as f:
            quota, period = f.read().split()
        if quota != "max":
            return max(1, round(int(quota) / int(period)))
    except Exception:
        pass
    # cgroup v1
    try:
        with open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") as f:
            quota = int(f.read())
        with open("/sys/fs/cgroup/cpu/cpu.cfs_period_us") as f:
            period = int(f.read())
        if quota > 0:
            return max(1, quota // period)
    except Exception:
        pass
    return os.cpu_count() or 2

_n_threads = _container_cpu_count()
print(f"Container CPU count: {_n_threads}  (os.cpu_count={os.cpu_count()})")

# ── Load at startup ───────────────────────────────────────────────────────────

if HF_TOKEN:
    login(token=HF_TOKEN)

_config_path = hf_hub_download(MODEL_REPO, "config.json", token=HF_TOKEN)
_int8_path   = hf_hub_download(MODEL_REPO, MODEL_FILE, token=HF_TOKEN)
# External-data sidecar ("<model>.onnx.data") exists ONLY when the model was
# exported with external tensors (models over the 2 GB protobuf limit). An int8
# model is usually self-contained and has no sidecar — fetch it only if present,
# so a self-contained model doesn't 404 the Space at startup. When it does exist,
# ONNX Runtime resolves it relative to the .onnx file; both land in the same HF
# snapshot dir, so co-locating them here is enough.
try:
    hf_hub_download(MODEL_REPO, f"{MODEL_FILE}.data", token=HF_TOKEN)
except EntryNotFoundError:
    pass

with open(_config_path, encoding="utf-8") as fh:
    _config = json.load(fh)

_id2label  = {int(k): v for k, v in _config["id2label"].items()}
_tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, subfolder="tokenizer", token=HF_TOKEN)

_opts = ort.SessionOptions()
_opts.intra_op_num_threads = _n_threads
_opts.inter_op_num_threads = 1
_opts.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
_sess = ort.InferenceSession(_int8_path, sess_options=_opts, providers=["CPUExecutionProvider"])

print(f"ONNX INT8 session ready — {_n_threads} threads, {len(_id2label)} classes")


# ── Inference ─────────────────────────────────────────────────────────────────

def _suppress_same_label(entities: list) -> list:
    kept = []
    for ent in sorted(entities, key=lambda x: -x["score"]):
        if any(k["entity_group"] == ent["entity_group"]
               and ent["start"] < k["end"] and k["start"] < ent["end"] for k in kept):
            continue
        kept.append(ent)
    return kept


def detect(text: str, threshold: float = THRESHOLD):
    enc = _tokenizer(
        text, max_length=MAX_LEN, truncation=True, stride=STRIDE, padding=True,
        return_overflowing_tokens=True, return_offsets_mapping=True, return_tensors="np",
    )
    all_offsets = enc.pop("offset_mapping")
    enc.pop("overflow_to_sample_mapping", None)

    best: dict = {}
    for w in range(all_offsets.shape[0]):
        feed = {
            "input_ids":      enc["input_ids"][w:w + 1].astype(np.int64),
            "attention_mask": enc["attention_mask"][w:w + 1].astype(np.int64),
        }
        logits  = _sess.run(["logits"], feed)[0][0]   # [classes, L, L]
        offsets = all_offsets[w]

        cidx, sidx, eidx = np.where(logits > threshold)
        if len(cidx) == 0:
            continue

        scores = 1.0 / (1.0 + np.exp(-logits[cidx, sidx, eidx]))
        cs = offsets[sidx, 0]
        ce = offsets[eidx, 1]
        keep = cs < ce
        for cid, c0, c1, sc in zip(cidx[keep], cs[keep], ce[keep], scores[keep]):
            key = (int(cid), int(c0), int(c1))
            if key not in best or sc > best[key]["score"]:
                best[key] = {
                    "entity_group": _id2label[int(cid)],
                    "start": int(c0), "end": int(c1),
                    "score": float(sc),
                }

    doc_len = max(1, len(text))
    cands = [e for e in best.values() if (e["end"] - e["start"]) <= 0.5 * doc_len]
    cands = _suppress_same_label(cands)

    flat = []
    for ent in sorted(cands, key=lambda x: (-(x["end"] - x["start"]), -x["score"])):
        if any(ent["start"] < k["end"] and k["start"] < ent["end"] for k in flat):
            continue
        flat.append(ent)
    flat = sorted(flat, key=lambda x: x["start"])
    return flat, cands


# ── UI helpers ────────────────────────────────────────────────────────────────

COLOR_MAP = {
    "PERSON_NAME": "#FF6B6B", "AADHAAR": "#6a0dad", "PASSPORT_NUMBER": "#1e90ff",
    "PAN": "#ff8c00", "SSN": "#dc143c", "DRIVER_LICENSE": "#ff7f50", "VOTER_ID": "#ff69b4",
    "PHONE_NUMBER": "#20b2aa", "EMAIL": "#2e8b57", "USERNAME": "#3cb371", "PASSWORD": "#8b0000",
    "BANK_ACCOUNT": "#8b4513", "CREDIT_CARD": "#ff1493", "CVV": "#c71585",
    "CREDIT_CARD_EXPIRATION": "#db7093", "UPI_ID": "#4682b4", "MONEY": "#2f4f4f",
    "DATE_OF_BIRTH": "#228b22", "DATE": "#6b8e23", "ADDRESS": "#8b008b",
    "LOCATION": "#008b8b", "ORGANIZATION": "#cd853f", "EMAIL_ADDRESS": "#2e8b57",
    "IP_ADDRESS": "#483d8b", "URL": "#4169e1", "API_KEY": "#800000",
}
_FALLBACK = ["#45B7D1", "#96CEB4", "#FFEAA7", "#DDA0DD", "#F7DC6F", "#85C1E9"]
_extra: dict = {}

def _label_color(label: str) -> str:
    if label in COLOR_MAP:
        return COLOR_MAP[label]
    if label not in _extra:
        _extra[label] = _FALLBACK[len(_extra) % len(_FALLBACK)]
    return _extra[label]


def _highlight_html(text: str, cands: list) -> str:
    if not cands:
        return f"<div style='line-height:2.4;font-size:15px;font-family:sans-serif'>{_html.escape(text)}</div>"

    ents = sorted(cands, key=lambda e: (e["start"], -(e["end"] - e["start"])))
    for e in ents:
        e["_children"] = []

    roots = []
    for e in ents:
        parent = None
        for c in ents:
            if (c is not e
                    and c["start"] <= e["start"] and e["end"] <= c["end"]
                    and (c["end"] - c["start"]) > (e["end"] - e["start"])):
                if parent is None or (c["end"] - c["start"]) < (parent["end"] - parent["start"]):
                    parent = c
        (parent["_children"] if parent else roots).append(e)

    def render(lo, hi, children):
        children = sorted(children, key=lambda c: c["start"])
        out, cur = "", lo
        for ch in children:
            if ch["start"] < cur:
                continue
            out  += _html.escape(text[cur:ch["start"]])
            color = _label_color(ch["entity_group"])
            label = ch["entity_group"]
            pct   = round(ch["score"] * 100, 1)
            inner = render(ch["start"], ch["end"], ch["_children"])
            badge = f'<small style="opacity:.8;font-size:10px"> {label}</small>'
            title = f"{label} {pct}%"
            out  += (f'<span style="background:{color};color:white;padding:2px 5px;'
                     f'border-radius:4px;font-weight:600;" title="{title}">'
                     f'{inner}{badge}</span>')
            cur   = ch["end"]
        out += _html.escape(text[cur:hi])
        return out

    body = render(0, len(text), roots)
    return f"<div style='line-height:2.6;font-size:15px;font-family:sans-serif'>{body}</div>"


def run_analyze(text: str, threshold: float):
    if not text or not text.strip():
        return "<div style='color:#888'>Paste text above and click Analyze.</div>", pd.DataFrame(), "", ""

    t0 = time.time()
    flat, cands = detect(text, threshold)
    elapsed = round(time.time() - t0, 2)

    html_out = _highlight_html(text, cands)

    redacted = text
    for e in sorted(flat, key=lambda x: x["start"], reverse=True):
        redacted = redacted[:e["start"]] + f'[{e["entity_group"]}]' + redacted[e["end"]:]

    rows = [{
        "Type":       e["entity_group"],
        "Text":       text[e["start"]:e["end"]].strip(),
        "Confidence": f'{round(e["score"] * 100, 1)}%',
        "Start":      e["start"],
        "End":        e["end"],
    } for e in flat]
    df = pd.DataFrame(rows)

    summary = f"{elapsed}s  |  {len(flat)} entities  |  {len(text)} chars  |  {_n_threads} CPU threads"
    return html_out, df, redacted, summary


# ── Gradio UI ─────────────────────────────────────────────────────────────────

EXAMPLES = [
    ["My name is John Smith, email john@example.com, SSN 123-45-6789, card 4532 8812 7731 9823.", 0.0],
    ["Patient Priya Sharma, DOB 12/03/1991, Aadhaar 2345 6789 0123, prescribed metformin 500mg.", 0.0],
    ["Invoice to Acme Corp, 400 Market St, San Francisco CA 94105. Contact: +1 415 555 0100.", 0.0],
]

with gr.Blocks(title="CLOAK PII Detector", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "## CLOAK V3.10 — PII Detector\n"
        "Detects personally identifiable information using nested NER (ONNX INT8 CPU). "
        "Nested highlighting shows overlapping spans."
    )

    with gr.Row():
        with gr.Column(scale=3):
            txt_input = gr.Textbox(
                label="Input text",
                placeholder="Paste text here…",
                lines=8, max_lines=30,
            )
            threshold_slider = gr.Slider(
                minimum=-2.0, maximum=6.0, value=0.0, step=0.5,
                label="Detection threshold  (higher = fewer, more confident detections)",
            )
            with gr.Row():
                btn       = gr.Button("Analyze", variant="primary", scale=2)
                clear_btn = gr.Button("Clear", scale=1)
            summary_box = gr.Textbox(label="", lines=1, interactive=False, show_label=False)

    gr.Examples(examples=EXAMPLES, inputs=[txt_input, threshold_slider], label="Try an example")

    highlighted = gr.HTML(label="Highlighted output")

    with gr.Row():
        with gr.Column():
            entity_table = gr.Dataframe(
                headers=["Type", "Text", "Confidence", "Start", "End"],
                label="Detected entities",
                interactive=False, wrap=True,
            )
        with gr.Column():
            redacted_out = gr.Textbox(label="Redacted text", lines=8, interactive=False)

    btn.click(
        fn=run_analyze,
        inputs=[txt_input, threshold_slider],
        outputs=[highlighted, entity_table, redacted_out, summary_box],
    )
    txt_input.submit(
        fn=run_analyze,
        inputs=[txt_input, threshold_slider],
        outputs=[highlighted, entity_table, redacted_out, summary_box],
    )
    clear_btn.click(
        fn=lambda: ("", pd.DataFrame(), "", ""),
        outputs=[txt_input, entity_table, redacted_out, summary_box],
    )

# ── REST API for the CLOAK backend (callHfSpace contract) ─────────────────────

api = FastAPI(title="CLOAK PII API", version="1.0.0")


class DetectRequest(BaseModel):
    text: str = Field(..., min_length=1, max_length=20_000)


def _fmt(text: str, e: dict) -> dict:
    """Shape one entity to the contract the CLOAK API expects."""
    return {
        "type":       e["entity_group"],
        "text":       text[e["start"]:e["end"]].strip(),
        "start":      e["start"],
        "end":        e["end"],
        "confidence": round(e["score"], 4),
    }


def _verify(request: Request):
    auth = request.headers.get("Authorization", "")
    if not auth.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing Authorization header")
    token = auth.removeprefix("Bearer ").strip()
    if not API_SECRET or token != API_SECRET:
        raise HTTPException(status_code=401, detail="Invalid token")


@api.get("/health")
def health():
    return {"status": "ok", "model": MODEL_REPO, "engine": "onnx-int8-cpu",
            "threads": _n_threads, "classes": len(_id2label)}


@api.post("/detect")
def detect_api(body: DetectRequest, request: Request):
    _verify(request)
    t0 = time.time()
    try:
        flat, cands = detect(body.text)
    except Exception as exc:
        raise HTTPException(status_code=503, detail=f"Inference failed: {exc}")

    entities = [_fmt(body.text, e) for e in flat
                if body.text[e["start"]:e["end"]].strip()]
    nested = [_fmt(body.text, e)
              for e in sorted(cands, key=lambda x: (x["start"], -(x["end"] - x["start"])))
              if body.text[e["start"]:e["end"]].strip()]

    redacted = body.text
    for e in sorted(flat, key=lambda x: x["start"], reverse=True):
        redacted = redacted[:e["start"]] + f'[{e["entity_group"]}]' + redacted[e["end"]:]

    return {
        "entities":           entities,
        "nested":             nested,
        "redacted_text":      redacted,
        "model":              MODEL_REPO,
        "processing_time_ms": int((time.time() - t0) * 1000),
    }


# Mount the Gradio UI at "/" on the same FastAPI app, so one Space serves both
# the demo UI and the /detect + /health API.
app = gr.mount_gradio_app(api, demo, path="/")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "7860")))