"""
Stage 2 - asking MiMo.

The prompt, the prefill and the parser are **lifted verbatim from Part B**
(`Evaluation_of_OpenSource_Models_for_PDF_Injection_Recognition/Final_Project_Evaluation_V1.ipynb`,
cell 48). Same system message, same closed family list in the same order, same MiMo prefill, same
brace-counting parser, same greedy decoding, same 200-token cap. A reworded prompt or a widened
parser is a different experiment, and Part B's measured F1 of 0.945 would no longer describe it.

**Two runtimes, chosen at startup, because a Space's hardware is not this code's decision.**

  - `gpu` - the original BF16 checkpoint quantised to 4-bit NF4 by `bitsandbytes`, batched exactly
    as Part B batched it. This *is* Part B's configuration; nothing about the arithmetic differs.
  - `cpu` - a Q4_K_M GGUF through `llama.cpp`, for a Space with no GPU at all. Same base model,
    different quantisation, and from a build with MiMo's multi-token-prediction layers stripped.

Which one is live is reported in the interface, with `RUNTIME_CAVEAT` next to it, because the two
do not deserve the same confidence in the published numbers.
"""

import json
import os
import re
import threading

FAMILIES = ["cross_site_scripting", "dde_template_injection", "javascript_injection",
            "llm_prompt_injection", "object_action_injection", "polyglot_file",
            "ransomware_simulation", "shellcode_embedded_exe", "ssrf",
            "steganographic_payload", "uri_redirect_phishing", "xfa_acroform_injection"]

# ---------------------------------------------------------------------------------------------
# Verbatim from Part B, cell 48
# ---------------------------------------------------------------------------------------------

SYSTEM = (
    "You are a PDF security analyst. You are given the raw extracted text of a PDF file - object "
    "definitions, stream contents and metadata, exactly as they appear in the file. Some of these "
    "files have had a malicious payload injected into them; most, but not all, have. Your job is to "
    "say which, and to point at your evidence.\n\n"
    "Answer with a single JSON object and nothing else:\n"
    '{"injected": true or false, '
    '"injection_type": one of ' + json.dumps(FAMILIES) + ' or "none", '
    '"evidence": the exact substring from the input that convinced you, at most 200 characters, '
    'or "" if none, '
    '"reasoning": one short sentence}\n\n'
    "If the file looks clean, answer injected=false and injection_type=\"none\". Do not guess a "
    "family when you do not believe there is an injection."
)

MAX_CHARS = 3000        # payload_window is already capped at this; belt and braces
MAX_NEW = 200
BATCH = 8               # Part B's batch, quartered automatically on CUDA OOM


def build_messages(text: str):
    """The chat turns for one window."""
    return [{"role": "system", "content": SYSTEM},
            {"role": "user", "content": "PDF extract:\n\n```\n" + text[:MAX_CHARS] + "\n```"}]


# Text appended to the assistant turn, so the model resumes from it instead of starting free.
# MiMo is reasoning-trained, opens every answer with `<think>`, and at MAX_NEW = 200 the budget is
# gone before the block closes - not one of its 1,100 Part B answers contained a closing
# `</think>`, so it never reached the JSON. An empty, already-closed block says the deliberation is
# finished before it begins, and the opening brace puts it inside the answer.
PREFILL = '<think>\n\n</think>\n\n{"injected":'


def scan_objects(raw: str):
    r"""
    Every balanced {...} in the text, counting braces and skipping anything inside a string literal.

    A regex cannot do this - matching balanced delimiters is outside what regular expressions can
    express - and the naive `\{[^{}]*\}` this replaced was actively harmful: every injected file in
    this corpus carries an EICAR-style marker containing a `}`, so the moment a model quoted its
    evidence the match was truncated mid-string and the verdict was discarded. The bug fired
    exactly when the model was RIGHT.
    """
    objs, depth, start, in_str, esc = [], 0, None, False, False
    for i, ch in enumerate(raw or ""):
        if in_str:
            if esc:
                esc = False
            elif ch == "\\":
                esc = True
            elif ch == '"':
                in_str = False
            continue
        if ch == '"':
            in_str = True
        elif ch == "{":
            if depth == 0:
                start = i
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0 and start is not None:
                objs.append(raw[start:i + 1])
                start = None
            depth = max(depth, 0)
    return objs


def parse_response(raw: str) -> dict:
    """
    Pull the verdict out of whatever the model said.

    The LAST balanced object is taken, not the first: reasoning-trained models restate the schema
    while thinking and emit the real answer at the end. If nothing parses as JSON the two fields
    that matter are lifted out individually rather than thrown away. Only a response with no
    recoverable verdict counts as parse_ok=False, and that reads as 'not injected': a detector that
    cannot make itself understood has caught nothing.

    Kept deliberately narrow, exactly as Part B scored it. A wider salvage would recover more
    verdicts and would also mean the F1 quoted in the UI describes a parser that is not this one.
    """
    for m in reversed(scan_objects(raw)):
        try:
            obj = json.loads(m)
        except json.JSONDecodeError:
            continue
        if "injected" in obj:
            inj = obj["injected"]
            inj = inj if isinstance(inj, bool) else str(inj).strip().lower() in {"true", "yes", "1"}
            fam = str(obj.get("injection_type", "none") or "none").strip().lower()
            return {"parse_ok": True, "parsed_by": "balanced JSON",
                    "pred_injected": int(inj),
                    "pred_family": fam if fam in FAMILIES else "none",
                    "evidence": str(obj.get("evidence", ""))[:200],
                    "reasoning": str(obj.get("reasoning", ""))[:300]}

    m = re.search(r'"injected"\s*:\s*(true|false)', raw or "", re.I)
    if m:
        f = re.search(r'"injection_type"\s*:\s*"([a-z_]+)"', raw, re.I)
        fam = f.group(1).lower() if f else "none"
        return {"parse_ok": True, "parsed_by": "field regex",
                "pred_injected": int(m.group(1).lower() == "true"),
                "pred_family": fam if fam in FAMILIES else "none",
                "evidence": "", "reasoning": ""}

    return {"parse_ok": False, "parsed_by": "unrecoverable", "pred_injected": 0,
            "pred_family": "none", "evidence": "", "reasoning": ""}


# ---------------------------------------------------------------------------------------------
# Which runtime is live
# ---------------------------------------------------------------------------------------------

BASE_REPO = "XiaomiMiMo/MiMo-7B-RL"

# MiMo carries Multi-Token-Prediction layers that llama.cpp cannot load, so the only GGUF that runs
# at all is one with those layers removed. MTP is a speculative-decoding accelerator - the ordinary
# forward pass does not use it - so greedy output should be unaffected, but "should be" is doing
# real work in that sentence, and the quantiser's own model card says the same.
GGUF_REPO = "quantflex/MiMo-7B-RL-nomtp-GGUF"
GGUF_FILE = "MiMo-7B-RL-nomtp-Q4_K_M.gguf"

# HF sets this on a ZeroGPU Space. It is checked instead of `torch.cuda.is_available()` because on
# ZeroGPU there is no device at import time - one is attached only inside a @spaces.GPU call - so
# asking torch at startup would answer "no GPU" on the very hardware that has one.
ZERO_GPU = bool(os.environ.get("SPACES_ZERO_GPU"))


def _cuda_present() -> bool:
    try:
        import torch
        return torch.cuda.is_available()
    except Exception:
        return False


BACKEND = "gpu" if (ZERO_GPU or _cuda_present()) else "cpu"


def _llama_installed() -> bool:
    """
    Whether the CPU runtime can actually run.

    On a Hugging Face Space it cannot, and that is not a bug to fix here: the prebuilt
    `llama-cpp-python` wheels are musl-linked while a Space is glibc, and compiling the sdist
    exceeds the build timeout. `requirements.txt` records both attempts. This flag lets the app
    stop offering a runtime that would only fail on being chosen.
    """
    import importlib.util
    return importlib.util.find_spec("llama_cpp") is not None


CPU_AVAILABLE = _llama_installed()
BACKENDS = [b for b in ("gpu", "cpu")
            if (b == "gpu" and BACKEND == "gpu") or (b == "cpu" and CPU_AVAILABLE)] or [BACKEND]

CAVEATS = {
    "gpu": ("This Space is running Part B's own configuration: the BF16 checkpoint of "
            "XiaomiMiMo/MiMo-7B-RL quantised to 4-bit NF4 by bitsandbytes, greedy, 200 new tokens, "
            "batched at 8. The prompt, prefill, decoding and parser are byte-identical to the run "
            "that produced F1 0.945, so that figure describes this configuration - measured on the "
            "project's own synthetic corpus, which is the thing it does not describe."),
    "cpu": ("This Space has no GPU, and bitsandbytes needs CUDA, so MiMo is running as a Q4_K_M "
            "GGUF through llama.cpp instead - from a build with the multi-token-prediction layers "
            "stripped, because llama.cpp cannot load them. The prompt, prefill, decoding and "
            "parser are byte-identical to Part B; the arithmetic underneath is not. Treat F1 0.945 "
            "as the figure for the configuration Part B measured, not for this one."),
}
RUNTIME_CAVEAT = CAVEATS[BACKEND]

# Measured in Part B on a T4 for the GPU path; measured loosely on 2 vCPUs for the CPU path.
SECONDS = {"gpu": 4.2, "cpu": 130}
SECONDS_PER_WINDOW = SECONDS[BACKEND]

_lock = threading.Lock()        # one model, one scan at a time
_tokenizer = None


def load_tokenizer():
    """
    MiMo's own tokenizer.

    On the GPU path it does the real work. On the CPU path it is used **only** to apply the chat
    template - the wrapper of role tags around the prompt - because Part B let each model's own
    tokenizer write that, and doing the same is what keeps the rendered string identical to the one
    that was measured. Only the tokenizer files are downloaded there, a few megabytes.
    """
    global _tokenizer
    if _tokenizer is None:
        from transformers import AutoTokenizer
        _tokenizer = AutoTokenizer.from_pretrained(BASE_REPO, trust_remote_code=True)
    return _tokenizer


def render_prompt(text: str):
    """
    Apply MiMo's chat template to one window, then append the prefill.

    If the tokenizer cannot be reached the ChatML fallback is used. MiMo is a ChatML model, so this
    produces the same string in practice - but it is a reconstruction rather than the model's own
    template, so the caller is told which route was taken instead of the difference being silent.
    """
    messages = build_messages(text)
    try:
        rendered = load_tokenizer().apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True)
        route = "tokenizer template"
    except Exception:
        rendered = "".join(f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>\n"
                           for m in messages) + "<|im_start|>assistant\n"
        route = "ChatML fallback"
    return rendered + PREFILL, route


# ---------------------------------------------------------------------------------------------
# GPU path - Part B's configuration, unchanged
# ---------------------------------------------------------------------------------------------

_gpu_model = None


def _gpu_decorator(duration: int):
    """
    `@spaces.GPU` on a ZeroGPU Space, and a no-op anywhere else.

    ZeroGPU scans for a decorated function **at startup** and refuses to serve the Space if it
    finds none, so the decoration has to happen at import time - which is why this is a decorator
    factory applied below rather than a check made when a scan begins.
    """
    try:
        import spaces
        return spaces.GPU(duration=duration)
    except Exception:
        return lambda fn: fn


def prefetch_weights(progress=None):
    """
    Pull the checkpoint to local disk **before** any GPU is requested.

    ZeroGPU bills wall-clock inside the decorated call and caps how long one may last, so a 15.7 GB
    first-run download in there does not merely waste the grant - it guarantees the very first scan
    exceeds the cap and fails. Downloading out here costs nothing but patience, and every later
    call finds the files cached.
    """
    from huggingface_hub import snapshot_download
    if progress:
        progress("fetching MiMo-7B weights (15.7 GB, first run only)")
    snapshot_download(BASE_REPO, allow_patterns=["*.safetensors", "*.json", "*.txt", "*.model"])


def _load_gpu():
    """MiMo in 4-bit NF4 on the attached device. Cached; ~5 GB of weights."""
    global _gpu_model
    if _gpu_model is None:
        import torch
        from transformers import AutoModelForCausalLM, BitsAndBytesConfig
        q = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                               bnb_4bit_compute_dtype=torch.float16,
                               bnb_4bit_use_double_quant=True)
        # device_map={"": 0} rather than "auto": everything on the one GPU or nothing. "auto" is
        # free to spill layers onto the CPU, which is exactly the state bitsandbytes then refuses
        # to work in, and the failure arrives as an offload complaint instead of "it does not fit".
        _gpu_model = AutoModelForCausalLM.from_pretrained(
            BASE_REPO, quantization_config=q, device_map={"": 0}, trust_remote_code=True)
        _gpu_model.eval()
    return _gpu_model


# Two separate limits squeeze this number, and the free one is much tighter than the hard one.
#
# The hard limit: the `spaces` client asks the scheduler for 1.5x whatever is set here, and ZeroGPU
# refuses anything over 300s - so above 200 is rejected outright, which is how the first deploy of
# this path died.
#
# The limit that actually bites: a free Hugging Face account gets roughly **five minutes of ZeroGPU
# per day**, and the scheduler reserves the full requested duration up front rather than what the
# run turns out to need. Asking for 180 (=270s reserved) therefore spends a whole day's quota on
# one batch. 110 covers the 4-bit load (~60s from cached weights) plus MAX_WINDOWS_GPU regions at
# ~4.2s each, reserves 165s, and leaves room for a second run in the same day.
GPU_DURATION = 110
MAX_WINDOWS_GPU = 8


@_gpu_decorator(GPU_DURATION)
def _judge_all_gpu(texts: list) -> list:
    """
    Every window in one GPU call, batched the way Part B batched documents.

    One call rather than one per window because ZeroGPU grants and reclaims the device around each
    decorated call, and paying the model-load cost per window would dominate everything else.
    """
    import torch

    model = _load_gpu()
    tok = load_tokenizer()
    rendered = [render_prompt(t) for t in texts]
    prompts = [p for p, _ in rendered]
    route = rendered[0][1] if rendered else "tokenizer template"

    tok.padding_side = "left"      # decoder-only: right padding starts generation after it
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token

    out, i, batch = [], 0, BATCH
    while i < len(prompts):
        chunk = prompts[i:i + batch]
        try:
            enc = tok(chunk, return_tensors="pt", padding=True,
                      truncation=True, max_length=2048).to(model.device)
            with torch.inference_mode():
                gen = model.generate(**enc, max_new_tokens=MAX_NEW, do_sample=False,
                                     pad_token_id=tok.pad_token_id)
            decoded = tok.batch_decode(gen[:, enc["input_ids"].shape[1]:], skip_special_tokens=True)
        except torch.cuda.OutOfMemoryError:
            torch.cuda.empty_cache()
            if batch == 1:
                raise
            batch = max(1, batch // 4)          # Part B's policy: shrink rather than hard-code
            continue

        # generate() returns only the new tokens, so the prefill was given to the model but never
        # generated. Put it back before parsing or the opening brace of the JSON is missing.
        for text in decoded:
            raw = PREFILL + text
            out.append({**parse_response(raw), "raw": raw.strip()[:2000], "prompt_route": route})
        i += len(chunk)
    return out


# ---------------------------------------------------------------------------------------------
# CPU path - llama.cpp
# ---------------------------------------------------------------------------------------------

N_CTX = 4096                     # a 3,000-char window is ~1,000 tokens, plus 200 generated
_llm = None


def _load_cpu(progress=None):
    """Load the GGUF once and keep it. ~4.7 GB resident, well inside a free tier's 16 GB."""
    global _llm
    if _llm is None:
        if not CPU_AVAILABLE:
            raise RuntimeError(
                "the CPU runtime needs llama-cpp-python, which is not installed on this Space. "
                "The prebuilt wheels are musl-linked (a Space is glibc) and building the sdist "
                "exceeds the Space build timeout, so only the GPU runtime is available here - "
                "which is also the one that reproduces Part B exactly.")
        from huggingface_hub import hf_hub_download
        from llama_cpp import Llama
        if progress:
            progress(f"downloading {GGUF_FILE} (4.7 GB, first run only)")
        path = hf_hub_download(GGUF_REPO, GGUF_FILE)
        if progress:
            progress("loading MiMo-7B into memory")
        _llm = Llama(model_path=path, n_ctx=N_CTX, n_threads=os.cpu_count() or 2,
                     n_batch=256, logits_all=False, verbose=False)
    return _llm


def _judge_one_cpu(text: str, progress=None) -> dict:
    """One window. Greedy, 200 new tokens, exactly as Part B decoded."""
    llm = _load_cpu(progress)
    prompt, route = render_prompt(text)
    out = llm.create_completion(prompt=prompt, max_tokens=MAX_NEW, temperature=0.0,
                                top_k=1, stop=["<|im_end|>", "<|endoftext|>"])
    raw = PREFILL + out["choices"][0]["text"]
    return {**parse_response(raw), "raw": raw.strip()[:2000], "prompt_route": route}


# ---------------------------------------------------------------------------------------------
# The one entry point the app uses
# ---------------------------------------------------------------------------------------------

def judge_all(texts: list, progress=None, backend: str = None) -> list:
    """
    Windows in, one parsed verdict each, in order.

    `backend` overrides the startup choice, which matters on ZeroGPU: a free account gets about
    five minutes of GPU per day, and once that is gone the Space is not broken - it still has a
    CPU. Running the llama.cpp path there is slow and quota-free, so the choice is offered rather
    than the whole app going dark until midnight. Asking for "gpu" on hardware that has none is
    refused here instead of failing somewhere inside bitsandbytes.

    The GPU path answers everything in a single call and so cannot report progress part-way; the
    CPU path is slow enough that per-window progress is the difference between a usable page and a
    frozen one. Both return the same shape.
    """
    if not texts:
        return []
    chosen = backend or BACKEND
    if chosen == "gpu" and BACKEND != "gpu":
        raise RuntimeError("this Space has no GPU; 4-bit loading through bitsandbytes needs CUDA")

    with _lock:
        if chosen == "gpu":
            prefetch_weights(progress)          # outside the GPU grant, deliberately
            if progress:
                progress(f"MiMo reading {len(texts)} region(s) on the GPU")
            return _judge_all_gpu(texts)
        out = []
        for i, t in enumerate(texts):
            if progress:
                progress(f"MiMo reading region {i + 1} of {len(texts)} on the CPU")
            out.append(_judge_one_cpu(t, progress))
        return out