import json
import re
from pathlib import Path

import gradio as gr
import torch
from transformers import AutoTokenizer
from granite_switch.hf import GraniteSwitchForCausalLM

MODEL_ID = "barha/granite-switch-4.0-350m-cti"
ADAPTER_NAME = "cti-technique-mapping"

# The adapter was trained on CTI sentences wrapped in this exact instruction
# format (see the ETL: USER_PROMPT + the sentence inside <cti>...</cti> tags).
# Sending bare CTI text instead makes the model fall back to base behavior and
# emit garbage (e.g. repeated <tool_call>), so we must reproduce the format here.
USER_PROMPT = "What ATT&CK technique does the following CTI procedure sentence describe?"

# MITRE ATT&CK Enterprise technique ID -> human-readable name. Built from the
# official mitre-attack/attack-stix-data Enterprise bundle (697 techniques +
# sub-techniques), bundled as a static file so the Space needs no network at
# inference and resolves any ID the adapter emits.
_ID_TO_NAME = json.loads((Path(__file__).parent / "mitre_id_to_name.json").read_text())

# Match a MITRE technique ID anywhere in a string, e.g. T1059 or T1059.001.
_TID_RE = re.compile(r"\bT\d{4}(?:\.\d{3})?\b")

# Load the model and tokenizer once at startup. The CTI adapter is activated by
# the <|cti-technique-mapping|> control token. The chat template only inserts
# that token when `adapter_name` is passed to apply_chat_template — the default
# render leaves it out, giving the plain base model. We run BOTH below.
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
# No device_map: the Space runs on single-device CPU (cpu-basic), and
# device_map="auto" crashes here because accelerate's check_device_map can't
# place the non-float `model.adapter_token_ids` buffer. Plain load → CPU.
model = GraniteSwitchForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32,
)
model.eval()

EXAMPLES = [
    "The actor used PowerShell to download and execute a payload from a remote server.",
    "The malware created a scheduled task to maintain persistence across reboots.",
    "The threat actor dumped credentials from LSASS memory using a custom tool.",
    "The adversary encrypted files on the victim host and dropped a ransom note.",
    "The implant communicated with its command-and-control server over HTTPS on port 443.",
]


def _generate(cti_text: str, use_adapter: bool, max_new_tokens: int) -> str:
    """Run one forward pass, with or without firing the CTI adapter."""
    user_content = f"{USER_PROMPT}\n\n<cti>\n{cti_text}\n</cti>"
    messages = [{"role": "user", "content": user_content}]
    # Passing adapter_name fires the <|cti-technique-mapping|> control token in
    # the chat template; omitting it renders the plain prompt → base model.
    template_kwargs = {"adapter_name": ADAPTER_NAME} if use_adapter else {}
    enc = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
        **template_kwargs,
    ).to(model.device)

    prompt_len = enc["input_ids"].shape[1]
    with torch.no_grad():
        out = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False)
    return tokenizer.decode(out[0, prompt_len:], skip_special_tokens=True).strip()


def _resolve_name(technique_id: str) -> str:
    """Look up the MITRE name for an ID; '' if not a known technique."""
    return _ID_TO_NAME.get(technique_id, "")


def compare(cti_text: str):
    """Run the base model and the adapter on the same input, side by side.

    Returns (base_output, adapter_output) markdown strings.
    """
    cti_text = (cti_text or "").strip()
    if not cti_text:
        msg = "_Enter some CTI text describing adversary behavior._"
        return msg, msg

    # Base model: no adapter. It tends to answer in prose, not a clean ID — give
    # it room to do so. This is the "before" half of the demo.
    base_raw = _generate(cti_text, use_adapter=False, max_new_tokens=48)
    base_md = base_raw or "_(no output)_"

    # Adapter: fires the control token; trained to emit exactly one technique ID.
    adapter_raw = _generate(cti_text, use_adapter=True, max_new_tokens=16)
    match = _TID_RE.search(adapter_raw)
    if match:
        tid = match.group(0)
        name = _resolve_name(tid)
        if name:
            adapter_md = f"### `{tid}`\n**{name}**"
        else:
            adapter_md = f"### `{tid}`\n_(name not in MITRE map)_"
    else:
        adapter_md = f"_(no technique ID returned)_\n\n```\n{adapter_raw}\n```"

    return base_md, adapter_md


with gr.Blocks(title="Granite Switch · CTI Technique Mapping") as demo:
    gr.Markdown(
        """
        # 🛡️ Granite Switch 4.0 350M — CTI Technique Mapping

        The **same base model** answers a **cyber threat intelligence (CTI)** prompt
        two ways. Without the adapter it rambles in prose; flip on the embedded
        **`cti-technique-mapping`** LoRA (via a single control token) and it emits the
        one matching **MITRE ATT&CK** technique ID — which we resolve to its name.

        Model: [`barha/granite-switch-4.0-350m-cti`](https://huggingface.co/barha/granite-switch-4.0-350m-cti)
        · Base: `ibm-granite/granite-4.0-350m` · Adapter: `cti-technique-mapping` (LoRA, attention + MLP)
        """
    )

    cti_input = gr.Textbox(
        label="CTI text",
        placeholder="Describe the observed adversary behavior…",
        lines=4,
    )
    submit = gr.Button("Compare base vs adapter", variant="primary")

    with gr.Row():
        with gr.Column():
            gr.Markdown("#### 🟡 Base model (adapter OFF)")
            base_out = gr.Markdown()
        with gr.Column():
            gr.Markdown("#### 🟢 Granite Switch (adapter ON)")
            adapter_out = gr.Markdown()

    gr.Examples(
        examples=[[e] for e in EXAMPLES],
        inputs=cti_input,
        outputs=[base_out, adapter_out],
        fn=compare,
        cache_examples=False,
    )

    submit.click(compare, inputs=cti_input, outputs=[base_out, adapter_out])
    cti_input.submit(compare, inputs=cti_input, outputs=[base_out, adapter_out])


if __name__ == "__main__":
    demo.launch()