"""
CLOAK model API: the Hugging Face Space the CLOAK server calls. API only, no UI.

One bearer secret (API_SECRET) guards every data path:

  POST /detect            personal data in text: nested NER, ONNX INT8 on CPU
                          {"text"} -> {entities, nested, redacted_text, model, processing_time_ms}
  POST /file/classify     what kind of document a text is            (docclass.py)
  POST /file/analyze      the same for raw file bytes, with OCR      (docclass.py)
  POST /copilot/parse     a policy described in plain English -> drafts (copilot_api.py)
  GET  /copilot/examples  sample requests for the dashboard

  GET  /health            status of all three, no secret needed and no secret revealed

Contracts are owned by the CLOAK server: server/src/detection/pii.ts (/detect, /health),
server/src/detection/classify.ts (/file/*), server/src/modules/org/routes/copilot.ts (/copilot/*).

Each path loads on its own. If the document model or the copilot cannot load, only its routes
answer 503 and /health says why; /detect keeps working. The PII model is the one hard
requirement: without it the Space does not start, which HF shows as a failed build.
"""

import hmac
import json
import logging
import os
import time

import numpy as np
import onnxruntime as ort
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse
from huggingface_hub import hf_hub_download
from pydantic import BaseModel, Field
from transformers import AutoTokenizer

import copilot_api

# 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

logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO").upper(),
                    format="%(asctime)s %(levelname)s %(name)s: %(message)s")
log = logging.getLogger("cloak.api")

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

MODEL_REPO = os.getenv("MODEL_REPO", "Wardline/CLOAK_V3.12_onnx")
MODEL_FILE = os.getenv("MODEL_FILE", "model.fused.int8.onnx")  # fusion+int8
HF_TOKEN   = os.getenv("HF_TOKEN")
# Bearer token the CLOAK server must send. Falls back to HF_TOKEN for old installs; set a
# dedicated random API_SECRET instead, so CLOAK's settings never hold the HF account token.
API_SECRET = os.getenv("API_SECRET") or HF_TOKEN

MAX_LEN   = 512   # total tokens per window, INCLUDING special tokens
STRIDE    = 128   # tokens shared between consecutive windows
THRESHOLD = 0.0

# Largest request body /detect will accept.
#
# This MUST be >= DETECT_MAX_CHARS in server/src/detection/pii.ts (80,000).
# It used to be 20,000, so every document between 20k and 80k characters came
# back 422 — and CLOAK treated that as transient, retrying four times over 14
# seconds and counting it toward its circuit breaker, which then paused PII
# detection for every other caller.
#
# The model is indifferent to this number: detect() slices the text into
# overlapping 512-token windows however long it is. This only bounds the body.
MAX_INPUT_CHARS = int(os.getenv("MAX_INPUT_CHARS", "80000"))

# ── 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()
log.info("container CPU count: %d (os.cpu_count=%s)", _n_threads, os.cpu_count())

# ── Load at startup ───────────────────────────────────────────────────────────
# The token is passed to each download, never written to disk with login().

_config_path = hf_hub_download(MODEL_REPO, "config.json", token=HF_TOKEN)
MODEL_DIR = os.getenv("MODEL_DIR", "/tmp/cloak_model")
_int8_path = hf_hub_download(MODEL_REPO, MODEL_FILE, token=HF_TOKEN, local_dir=MODEL_DIR)

# 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, so an unguarded download
# raises EntryNotFoundError at import time and the Space never starts at all.
# Fetch it if it is there; carry on if it is not. When it does exist, ONNX
# Runtime resolves it relative to the .onnx file — both land in MODEL_DIR, so
# co-locating them here is enough.
try:
    hf_hub_download(MODEL_REPO, f"{MODEL_FILE}.data", token=HF_TOKEN, local_dir=MODEL_DIR)
except EntryNotFoundError:
    log.info("no external-data sidecar for %s: the model is self-contained", MODEL_FILE)

log.info("model path: %s", _int8_path)

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"])

log.info("ONNX INT8 session ready: %d threads, %d classes", _n_threads, len(_id2label))


# ── Inference ─────────────────────────────────────────────────────────────────
# Sliding windows are built BY HAND on purpose. Relying on the tokenizer's
# `return_overflowing_tokens` is unsafe: tokenizers 0.23.x returns only the first
# window plus a tiny stub (~stride+4 tokens) and silently drops the rest of the
# text, so anything after ~token 512 was never scored (0.20-0.22 were fine).

def _special_token_layout():
    """Token ids the tokenizer wraps around a sequence, e.g. [CLS] ... [SEP]."""
    probe = _tokenizer("a", return_offsets_mapping=True)
    ids, offs = probe["input_ids"], probe["offset_mapping"]
    real = [i for i, (s, e) in enumerate(offs) if e > s]
    return list(ids[:real[0]]), list(ids[real[-1] + 1:])


_PREFIX_IDS, _SUFFIX_IDS = _special_token_layout()
_BODY = MAX_LEN - len(_PREFIX_IDS) - len(_SUFFIX_IDS)   # content tokens per window
# A plain `assert` disappears under `python -O`, and a non-positive step here
# makes _make_windows loop forever — so this is a real check, not an assertion.
if not 0 <= STRIDE < _BODY:
    raise ValueError(f"STRIDE ({STRIDE}) must be >= 0 and smaller than the window body ({_BODY})")
log.info("windowing: %d content tokens/window, stride %d, specials %s+%s",
         _BODY, STRIDE, _PREFIX_IDS, _SUFFIX_IDS)


def _make_windows(text: str):
    """Tokenize the WHOLE text once, then slice it into overlapping windows."""
    enc = _tokenizer(text, add_special_tokens=False,
                     return_offsets_mapping=True, truncation=False)
    ids  = list(enc["input_ids"])
    offs = [tuple(o) for o in enc["offset_mapping"]]
    n    = len(ids)
    step = _BODY - STRIDE
    windows, s = [], 0
    while True:
        e = min(s + _BODY, n)
        windows.append((s, e))
        if e >= n:
            break
        s += step
    return ids, offs, windows


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):
    ids, offs, windows = _make_windows(text)
    if not ids:
        return [], []
    log.debug("detect: %d tokens -> %d window(s)", len(ids), len(windows))

    n_pre, n_suf = len(_PREFIX_IDS), len(_SUFFIX_IDS)
    best: dict = {}
    for s, e in windows:
        # Each window is run at its real length: no padding, so no pad-token
        # artifacts in the last (shortest) window.
        w_ids   = _PREFIX_IDS + ids[s:e] + _SUFFIX_IDS
        offsets = np.array([(0, 0)] * n_pre + offs[s:e] + [(0, 0)] * n_suf, dtype=np.int64)
        input_ids = np.asarray([w_ids], dtype=np.int64)
        feed = {"input_ids": input_ids, "attention_mask": np.ones_like(input_ids)}
        logits = _sess.run(["logits"], feed)[0][0]   # [classes, L, L]

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

        # Only real (non-special) tokens, and start must not come after end.
        valid = offsets[:, 1] > offsets[:, 0]
        ok = valid[sidx] & valid[eidx] & (sidx <= eidx)
        cidx, sidx, eidx = cidx[ok], sidx[ok], eidx[ok]
        if len(cidx) == 0:
            continue

        scores = 1.0 / (1.0 + np.exp(-logits[cidx, sidx, eidx]))
        cs = offsets[sidx, 0]
        ce = offsets[eidx, 1]
        for cid, c0, c1, sc in zip(cidx, cs, ce, scores):
            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),
                }

    # Drop absurdly long spans, but never on short inputs (a lone email/name
    # legitimately covers ~100% of a short text).
    max_span = max(0.5 * len(text), 100)
    cands = [e for e in best.values() if (e["end"] - e["start"]) <= max_span]
    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


# ── HTTP API ──────────────────────────────────────────────────────────────────

# No interactive docs or schema on a public URL: they would map every route for anyone.
app = FastAPI(title="CLOAK model API", version="2.0.0", docs_url=None, redoc_url=None, openapi_url=None)

# Without a secret, _verify rejects EVERY request. The Space looks healthy (it boots, /health
# is green) while every call returns 401. Say it once, loudly, at startup instead.
if not API_SECRET:
    log.warning("neither API_SECRET nor HF_TOKEN is set: every data path will answer 401")
elif not os.getenv("API_SECRET"):
    log.warning("API_SECRET is not set: falling back to HF_TOKEN. Set a dedicated random API_SECRET.")
elif len(API_SECRET) < 32:
    log.warning("API_SECRET is shorter than 32 characters: use a long random value")


def _verify(request: Request) -> None:
    auth = request.headers.get("Authorization", "")
    if not auth.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing Authorization header")
    token = auth.removeprefix("Bearer ").strip()
    # compare_digest, not ==: a plain string comparison returns as soon as two
    # bytes differ, which leaks the secret's prefix to anyone timing the endpoint.
    if not API_SECRET or not hmac.compare_digest(token.encode(), API_SECRET.encode()):
        raise HTTPException(status_code=401, detail="Invalid token")


# Body limits, checked from Content-Length BEFORE the body is read, so an oversized request
# costs nothing. Paths not listed are bounded by their own validation (pydantic / docclass).
_BODY_LIMITS = {
    "/detect": MAX_INPUT_CHARS * 4 + 4_096,   # 4 bytes per character at worst, plus the JSON
}


@app.middleware("http")
async def _guard(request: Request, call_next):
    path = request.url.path
    limit = _BODY_LIMITS.get(path) or (copilot_api.MAX_BODY_BYTES if path.startswith("/copilot/") else None)
    if limit and request.method == "POST":
        length = request.headers.get("content-length")
        if length is None:
            return JSONResponse({"detail": "Content-Length required"}, status_code=411)
        if not length.isdigit() or int(length) > limit:
            return JSONResponse({"detail": f"Body larger than {limit} bytes"}, status_code=413)
    response = await call_next(request)
    # An API: nothing here may be cached, sniffed or framed.
    response.headers["Cache-Control"] = "no-store"
    response.headers["X-Content-Type-Options"] = "nosniff"
    response.headers["X-Frame-Options"] = "DENY"
    response.headers["Referrer-Policy"] = "no-referrer"
    return response


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


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),
    }


@app.get("/")
def root():
    return {"service": "CLOAK model API", "status": "ok"}


@app.get("/health")
def health():
    # max_input_chars and auth_configured are reported so a size mismatch or a
    # missing secret is visible from one GET, instead of only showing up as
    # failed scans.
    return {"status": "ok", "model": MODEL_REPO, "engine": "onnx-int8-cpu",
            "threads": _n_threads, "classes": len(_id2label),
            "max_input_chars": MAX_INPUT_CHARS,
            "auth_configured": bool(API_SECRET),
            "docclass": docclass.status(),
            "copilot": copilot_api.status()}


@app.post("/detect")
def detect_api(body: DetectRequest, request: Request):
    _verify(request)
    t0 = time.time()
    try:
        flat, cands = detect(body.text)
    except Exception:
        # The reason goes to the log, not to the caller.
        log.exception("inference failed")
        raise HTTPException(status_code=503, detail="Inference failed")

    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),
    }


# Second path: document classification (docclass.py). Loaded after the PII
# model and fully isolated from it — a failure here leaves /detect running.
import docclass  # noqa: E402

docclass.load(HF_TOKEN)
app.include_router(docclass.build_router(_verify))

# Third path: Policy Copilot (copilot_api.py). Pure Python, no model download; isolated the same way.
copilot_api.load()
app.include_router(copilot_api.build_router(_verify))


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "7860")),
                proxy_headers=True, forwarded_allow_ips="*", server_header=False)