"""SigmaForge training utility Space (ZeroGPU).

Fine-tunes a code LLM to generate and repair Sigma detection rules using the
SigmaForge dataset. Training runs in bounded ZeroGPU chunks: each chunk loads
the latest LoRA adapter from the Hub, trains on a dataset shard, and pushes
the updated adapter back. A final merge step publishes the merged model.
All heavy compute happens on Hugging Face hardware.
"""

import json
import os
import shutil
import tempfile
import traceback

import gradio as gr
import spaces
import torch
from datasets import concatenate_datasets, load_dataset
from huggingface_hub import HfApi, snapshot_download

BASE_MODEL = "Qwen/Qwen2.5-Coder-1.5B-Instruct"
DATASET_REPO = "alirezaaminzadeh/sigmaforge-detection-rules"
OUTPUT_REPO = "alirezaaminzadeh/sigmaforge-rule-generator"
ADAPTER_PATH = "adapter"
TOKEN = os.environ["HF_TOKEN"]
SEED = 20260809
MAX_LEN = 1024

api = HfApi(token=TOKEN)

SYSTEM_PROMPT = (
    "You are SigmaForge, an expert detection engineering assistant. "
    "You write and repair Sigma detection rules as valid Sigma YAML. "
    "Respond with the YAML rule only, no extra commentary."
)


def build_training_dataset():
    """Assemble chat-formatted samples from the three dataset configs."""
    samples = []

    d1 = load_dataset(DATASET_REPO, "description_to_sigma", split="train", token=TOKEN)
    for row in d1:
        hints = ", ".join(
            f"{k}={row[f'logsource_{k}']}"
            for k in ("product", "category", "service")
            if row[f"logsource_{k}"]
        )
        user = f"Write a Sigma rule for the following detection.\n\nDetection: {row['description']}"
        if hints:
            user += f"\nLog source: {hints}"
        if row["attack_techniques"]:
            user += f"\nMITRE ATT&CK: {', '.join(row['attack_techniques'])}"
        samples.append((user, row["sigma_rule"]))

    d3 = load_dataset(DATASET_REPO, "sigma_repair", split="train", token=TOKEN)
    for row in d3:
        user = (
            "The following Sigma rule fails validation. Fix it and return the corrected rule.\n\n"
            f"Rule:\n{row['broken_rule']}\n\nValidation errors:\n- "
            + "\n- ".join(row["validation_errors"])
        )
        samples.append((user, row["corrected_rule"]))

    d2 = load_dataset(DATASET_REPO, "logs_to_sigma", split="train", token=TOKEN)
    for i, row in enumerate(d2):
        if i >= 1200:
            break
        user = (
            "Write a Sigma rule that matches the positive events and not the negative events.\n\n"
            f"Positive events:\n{row['positive_events']}\n\nNegative events:\n{row['negative_events']}"
        )
        samples.append((user, row["expected_rule"]))

    return samples


def _state() -> dict:
    try:
        import requests

        url = f"https://huggingface.co/{OUTPUT_REPO}/raw/main/training_state.json"
        r = requests.get(url, headers={"Authorization": f"Bearer {TOKEN}"}, timeout=30)
        if r.status_code == 200:
            return r.json()
    except Exception:
        pass
    return {"chunks_done": 0, "samples_seen": 0}


def _save_state(state: dict) -> None:
    with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as f:
        json.dump(state, f, indent=2)
        path = f.name
    api.upload_file(path_or_fileobj=path, path_in_repo="training_state.json", repo_id=OUTPUT_REPO)


def _estimate_duration(shard_index: int, steps: int) -> int:
    # model load + tokenization overhead, plus per-step training cost
    return min(280, 100 + int(steps * 1.2))


@spaces.GPU(duration=_estimate_duration)
def _train_chunk_gpu(shard_index: int, steps: int):
    from peft import LoraConfig, PeftModel
    from transformers import AutoModelForCausalLM, AutoTokenizer
    from trl import SFTConfig, SFTTrainer
    from datasets import Dataset

    samples = build_training_dataset()
    ds = Dataset.from_list(
        [
            {
                "messages": [
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": u},
                    {"role": "assistant", "content": a},
                ]
            }
            for u, a in samples
        ]
    ).shuffle(seed=SEED)

    per_device = 4
    grad_accum = 2
    shard_size = steps * per_device * grad_accum
    start = shard_index * shard_size
    shard = ds.select(range(start % len(ds), min(start % len(ds) + shard_size, len(ds))))

    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
    model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16)

    peft_config = None
    try:
        adapter_dir = snapshot_download(OUTPUT_REPO, allow_patterns=[f"{ADAPTER_PATH}/*"], token=TOKEN)
        adapter_dir = os.path.join(adapter_dir, ADAPTER_PATH)
        if os.path.exists(os.path.join(adapter_dir, "adapter_config.json")):
            model = PeftModel.from_pretrained(model, adapter_dir, is_trainable=True)
            print("Resumed adapter from Hub")
        else:
            raise FileNotFoundError
    except Exception:
        from peft import LoraConfig

        peft_config = LoraConfig(
            r=16,
            lora_alpha=32,
            lora_dropout=0.05,
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        )
        print("Starting fresh LoRA adapter")

    out_dir = tempfile.mkdtemp()
    trainer = SFTTrainer(
        model=model,
        train_dataset=shard,
        peft_config=peft_config,
        processing_class=tokenizer,
        args=SFTConfig(
            output_dir=out_dir,
            max_length=MAX_LEN,
            per_device_train_batch_size=per_device,
            gradient_accumulation_steps=grad_accum,
            learning_rate=1e-4,
            lr_scheduler_type="constant",
            warmup_steps=10 if shard_index == 0 else 0,
            logging_steps=10,
            max_steps=steps,
            bf16=True,
            report_to=[],
            save_strategy="no",
            seed=SEED + shard_index,
        ),
    )
    trainer.train()
    loss = None
    for entry in reversed(trainer.state.log_history):
        if "loss" in entry or "train_loss" in entry:
            loss = entry.get("train_loss", entry.get("loss"))
            break

    save_dir = os.path.join(out_dir, "adapter_out")
    trainer.model.save_pretrained(save_dir)
    api.upload_folder(folder_path=save_dir, path_in_repo=ADAPTER_PATH, repo_id=OUTPUT_REPO)
    shutil.rmtree(out_dir, ignore_errors=True)
    return {"loss": loss, "trained_samples": len(shard), "total_samples": len(ds)}


def train_chunk(steps: float) -> str:
    try:
        steps = int(steps)
        state = _state()
        result = _train_chunk_gpu(state["chunks_done"], steps)
        state["chunks_done"] += 1
        state["samples_seen"] += result["trained_samples"]
        state["last_loss"] = result["loss"]
        state["total_samples"] = result["total_samples"]
        _save_state(state)
        return json.dumps({"ok": True, **state}, indent=2)
    except Exception:
        return json.dumps({"ok": False, "error": traceback.format_exc()})


@spaces.GPU(duration=180)
def _merge_gpu(out_dir: str) -> None:
    from peft import PeftModel
    from transformers import AutoModelForCausalLM, AutoTokenizer

    adapter_dir = snapshot_download(OUTPUT_REPO, allow_patterns=[f"{ADAPTER_PATH}/*"], token=TOKEN)
    adapter_dir = os.path.join(adapter_dir, ADAPTER_PATH)
    model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16).to("cuda")
    model = PeftModel.from_pretrained(model, adapter_dir)
    model = model.merge_and_unload()
    model = model.to("cpu")
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
    model.save_pretrained(out_dir)
    tokenizer.save_pretrained(out_dir)


def merge_and_push() -> str:
    """Merge the LoRA adapter into the base model (on GPU) and push."""
    try:
        out = tempfile.mkdtemp()
        _merge_gpu(out)
        api.upload_folder(folder_path=out, path_in_repo=".", repo_id=OUTPUT_REPO)
        shutil.rmtree(out, ignore_errors=True)
        return json.dumps({"ok": True, "merged_to": OUTPUT_REPO})
    except Exception:
        return json.dumps({"ok": False, "error": traceback.format_exc()})


def status() -> str:
    return json.dumps(_state(), indent=2)


# ---------------------------------------------------------------------------
# Evaluation
# ---------------------------------------------------------------------------

@spaces.GPU(duration=280)
def _eval_gpu(n_samples: int) -> dict:
    import re

    import yaml
    from transformers import AutoModelForCausalLM, AutoTokenizer

    from sigma.collection import SigmaCollection
    from sigma.backends.splunk import SplunkBackend
    from sigma.backends.elasticsearch import LuceneBackend

    tokenizer = AutoTokenizer.from_pretrained(OUTPUT_REPO)
    model = AutoModelForCausalLM.from_pretrained(OUTPUT_REPO, torch_dtype=torch.bfloat16).to("cuda")
    model.eval()

    test = load_dataset(DATASET_REPO, "description_to_sigma", split="test", token=TOKEN)
    test = test.shuffle(seed=SEED).select(range(min(n_samples, len(test))))

    backends = {"splunk": SplunkBackend(), "elastic": LuceneBackend()}
    stats = {
        "n": 0,
        "valid_yaml": 0,
        "schema_pass": 0,
        "compile_splunk": 0,
        "compile_elastic": 0,
        "attack_tp": 0,
        "attack_fp": 0,
        "attack_fn": 0,
    }

    def extract_yaml(text: str) -> str:
        m = re.search(r"```(?:yaml|yml)?\s*(.*?)```", text, re.DOTALL)
        return m.group(1).strip() if m else text.strip()

    for row in test:
        hints = ", ".join(
            f"{k}={row[f'logsource_{k}']}"
            for k in ("product", "category", "service")
            if row[f"logsource_{k}"]
        )
        user = f"Write a Sigma rule for the following detection.\n\nDetection: {row['description']}"
        if hints:
            user += f"\nLog source: {hints}"
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user},
        ]
        enc = tokenizer.apply_chat_template(
            messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
        ).to(model.device)
        with torch.no_grad():
            out = model.generate(
                **enc, max_new_tokens=600, do_sample=False, pad_token_id=tokenizer.eos_token_id
            )
        prompt_len = enc["input_ids"].shape[1]
        rule = extract_yaml(tokenizer.decode(out[0][prompt_len:], skip_special_tokens=True))
        stats["n"] += 1

        try:
            doc = yaml.safe_load(rule)
            assert isinstance(doc, dict)
            stats["valid_yaml"] += 1
        except Exception:
            continue

        schema_ok = all(k in doc for k in ("title", "logsource", "detection")) and isinstance(
            doc.get("detection"), dict
        ) and "condition" in doc["detection"]
        if schema_ok:
            stats["schema_pass"] += 1

        for name, backend in backends.items():
            try:
                backend.convert(SigmaCollection.from_yaml(rule))
                stats[f"compile_{name}"] += 1
            except Exception:
                pass

        pred = set()
        for tag in doc.get("tags") or []:
            m = re.match(r"attack\.(t\d{4}(?:\.\d{3})?)", str(tag), re.IGNORECASE)
            if m:
                pred.add(m.group(1).upper())
        gold = set(row["attack_techniques"])
        stats["attack_tp"] += len(pred & gold)
        stats["attack_fp"] += len(pred - gold)
        stats["attack_fn"] += len(gold - pred)

    n = max(stats["n"], 1)
    precision = stats["attack_tp"] / max(stats["attack_tp"] + stats["attack_fp"], 1)
    recall = stats["attack_tp"] / max(stats["attack_tp"] + stats["attack_fn"], 1)
    f1 = 2 * precision * recall / max(precision + recall, 1e-9)
    return {
        "samples": stats["n"],
        "valid_yaml_rate": round(stats["valid_yaml"] / n, 4),
        "schema_pass_rate": round(stats["schema_pass"] / n, 4),
        "splunk_compilation_rate": round(stats["compile_splunk"] / n, 4),
        "elastic_compilation_rate": round(stats["compile_elastic"] / n, 4),
        "attack_mapping_precision": round(precision, 4),
        "attack_mapping_recall": round(recall, 4),
        "attack_mapping_f1": round(f1, 4),
    }


def evaluate(n_samples: float) -> str:
    try:
        results = _eval_gpu(int(n_samples))
        with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as f:
            json.dump(results, f, indent=2)
            path = f.name
        api.upload_file(path_or_fileobj=path, path_in_repo="eval_results.json", repo_id=OUTPUT_REPO)
        return json.dumps({"ok": True, **results}, indent=2)
    except Exception:
        return json.dumps({"ok": False, "error": traceback.format_exc()})


with gr.Blocks(title="SigmaForge Trainer") as demo:
    gr.Markdown("# SigmaForge Trainer\nInternal training utility (ZeroGPU chunked SFT).")
    steps_in = gr.Number(value=150, label="Steps for this chunk")
    out = gr.Textbox(label="Result", lines=14)
    gr.Button("Train chunk").click(train_chunk, inputs=steps_in, outputs=out, api_name="train_chunk")
    gr.Button("Merge and push").click(merge_and_push, outputs=out, api_name="merge_and_push")
    gr.Button("Status").click(status, outputs=out, api_name="status")
    n_eval = gr.Number(value=60, label="Eval samples")
    gr.Button("Evaluate").click(evaluate, inputs=n_eval, outputs=out, api_name="evaluate")

demo.launch()