#!/usr/bin/env python3
"""
Bug Bounty Assistant Evaluation Suite
Evaluates Base Model vs Fine-Tuned Model on Test Split:
- CWE Accuracy
- CVSS Vector Exact Match & Score MAE
- Scope Check Accuracy
- JSON Validity & Field-Level F1 (Note Extraction)
- ROUGE-1 / ROUGE-2 / ROUGE-L (Report Writing & Improvement)
- Per-Language Breakdown (EN vs ID)
Saves: eval_results.json + generates comparison markdown table.
"""

import os
import sys
import json
import re
import math
import torch
import pandas as pd
from typing import Dict, List, Any
from rouge_score import rouge_scorer
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

SYSTEM_PROMPT = (
    "You are BugBounty-Triage-Assistant, an autonomous expert cybersecurity assistant specializing in "
    "web application security, bug bounty reporting, severity triage (CWE & CVSS v3.1), scope verification, "
    "and defensive remediation. Maintain absolute technical precision and clarity."
)

def extract_cwe(text: str) -> str:
    m = re.search(r"CWE-\d+", text)
    return m.group(0) if m else "None"

def extract_cvss(text: str) -> str:
    m = re.search(r"CVSS:3\.1/[A-Z0-9:/]+", text)
    return m.group(0) if m else "None"

def extract_score(text: str) -> float:
    m = re.search(r"(?:Score|score)[:\s]+(\d+\.\d+)", text)
    if m:
        try:
            return float(m.group(1))
        except ValueError:
            pass
    return -1.0

def extract_scope_verdict(text: str) -> str:
    lower = text.lower()
    if "out_of_scope" in lower or "out of scope" in lower or "di luar cakupan" in lower:
        return "out_of_scope"
    elif "needs_clarification" in lower or "needs clarification" in lower or "klarifikasi" in lower:
        return "needs_clarification"
    elif "in_scope" in lower or "in scope" in lower or "dalam cakupan" in lower:
        return "in_scope"
    return "unknown"

def evaluate_suite(data_dir: str, base_model_id: str, ft_model_dir: str, num_samples: int = 15):
    print("Initializing evaluation suite on test split...")
    test_path = os.path.join(data_dir, "all", "test.parquet")
    df = pd.read_parquet(test_path)

    scorer = rouge_scorer.RougeScorer(['rouge1', 'rouge2', 'rougeL'], use_stemmer=True)

    tokenizer = AutoTokenizer.from_pretrained(base_model_id)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    print(f"Loading Base Model: {base_model_id}...")
    base_model = AutoModelForCausalLM.from_pretrained(
        base_model_id,
        torch_dtype=torch.float32,
        low_cpu_mem_usage=True
    )
    base_model.eval()

    print(f"Loading Fine-Tuned Model: {ft_model_dir}...")
    ft_model = AutoModelForCausalLM.from_pretrained(
        ft_model_dir,
        torch_dtype=torch.float32,
        low_cpu_mem_usage=True
    )
    ft_model.eval()

    def generate_response(model, prompt_text: str) -> str:
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt_text}
        ]
        input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=40,
                do_sample=False,
                pad_token_id=tokenizer.eos_token_id
            )
        gen_tokens = outputs[0][inputs["input_ids"].shape[1]:]
        return tokenizer.decode(gen_tokens, skip_special_tokens=True).strip()

    tasks_to_test = ["severity_triage", "scope_check", "note_extraction", "report_improvement"]
    eval_data = {}
    for t in tasks_to_test:
        sub = df[df["task"] == t].sample(n=min(len(df[df["task"] == t]), num_samples), random_state=42)
        eval_data[t] = sub

    results = {
        "base_model": base_model_id,
        "finetuned_model": "BugBounty-Triage-Assistant-3B",
        "metrics": {
            "base": {},
            "finetuned": {}
        },
        "per_language": {
            "en": {"base": {}, "finetuned": {}},
            "id": {"base": {}, "finetuned": {}}
        }
    }

    print("Evaluating Severity Triage (CWE & CVSS)...")
    for model_name, model in [("base", base_model), ("finetuned", ft_model)]:
        cwe_correct = 0
        cvss_exact = 0
        score_diffs = []
        total = 0

        for _, r in eval_data["severity_triage"].iterrows():
            total += 1
            pred = generate_response(model, r["prompt"])
            target = r["response"]

            target_cwe = extract_cwe(target)
            pred_cwe = extract_cwe(pred)
            if pred_cwe == target_cwe and target_cwe != "None":
                cwe_correct += 1

            target_cvss = extract_cvss(target)
            pred_cvss = extract_cvss(pred)
            if pred_cvss == target_cvss and target_cvss != "None":
                cvss_exact += 1

            t_score = extract_score(target)
            p_score = extract_score(pred)
            if t_score >= 0 and p_score >= 0:
                score_diffs.append(abs(t_score - p_score))
            else:
                score_diffs.append(3.5) # Penalty for missing/invalid score

        results["metrics"][model_name]["cwe_accuracy"] = round(cwe_correct / total, 3)
        results["metrics"][model_name]["cvss_vector_exact_match"] = round(cvss_exact / total, 3)
        results["metrics"][model_name]["cvss_score_mae"] = round(sum(score_diffs) / len(score_diffs), 2)

    print("Evaluating Scope Check Accuracy...")
    for model_name, model in [("base", base_model), ("finetuned", ft_model)]:
        correct = 0
        total = 0
        for _, r in eval_data["scope_check"].iterrows():
            total += 1
            pred = generate_response(model, r["prompt"])
            target_v = extract_scope_verdict(r["response"])
            pred_v = extract_scope_verdict(pred)
            if pred_v == target_v and target_v != "unknown":
                correct += 1
        results["metrics"][model_name]["scope_check_accuracy"] = round(correct / total, 3)

    print("Evaluating Note Extraction (JSON validity & field F1)...")
    for model_name, model in [("base", base_model), ("finetuned", ft_model)]:
        valid_json = 0
        f1_scores = []
        total = 0
        for _, r in eval_data["note_extraction"].iterrows():
            total += 1
            pred = generate_response(model, r["prompt"])
            target_json = json.loads(r["response"])
            try:
                # Find JSON block
                m = re.search(r"\{.*\}", pred, re.DOTALL)
                pred_json = json.loads(m.group(0)) if m else json.loads(pred)
                valid_json += 1
                matched_keys = sum(1 for k in target_json if k in pred_json)
                precision = matched_keys / max(len(pred_json), 1)
                recall = matched_keys / len(target_json)
                f1 = (2 * precision * recall) / (precision + recall) if (precision + recall) > 0 else 0
                f1_scores.append(f1)
            except Exception:
                f1_scores.append(0.0)

        results["metrics"][model_name]["json_validity_rate"] = round(valid_json / total, 3)
        results["metrics"][model_name]["extraction_field_f1"] = round(sum(f1_scores) / total, 3)

    print("Evaluating Report Improvement ROUGE Scores...")
    for model_name, model in [("base", base_model), ("finetuned", ft_model)]:
        r1_list, r2_list, rl_list = [], [], []
        for _, r in eval_data["report_improvement"].iterrows():
            pred = generate_response(model, r["prompt"])
            target = r["response"]
            scores = scorer.score(target, pred)
            r1_list.append(scores["rouge1"].fmeasure)
            r2_list.append(scores["rouge2"].fmeasure)
            rl_list.append(scores["rougeL"].fmeasure)

        results["metrics"][model_name]["rouge_1"] = round(sum(r1_list) / len(r1_list), 3)
        results["metrics"][model_name]["rouge_2"] = round(sum(r2_list) / len(r2_list), 3)
        results["metrics"][model_name]["rouge_l"] = round(sum(rl_list) / len(rl_list), 3)

    # Per-language breakdown for report improvement
    for lang in ["en", "id"]:
        lang_sub = df[(df["task"] == "report_improvement") & (df["language"] == lang)].sample(n=1, random_state=42)
        for model_name, model in [("base", base_model), ("finetuned", ft_model)]:
            r_list = []
            for _, r in lang_sub.iterrows():
                pred = generate_response(model, r["prompt"])
                scores = scorer.score(r["response"], pred)
                r_list.append(scores["rougeL"].fmeasure)
            results["per_language"][lang][model_name]["rouge_l"] = round(sum(r_list) / len(r_list), 3)

    # Save to eval_results.json
    out_eval_path = os.path.join(os.path.dirname(ft_model_dir), "eval_results.json")
    with open(out_eval_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"Evaluation finished! Results saved to: {out_eval_path}")

    # Print comparison table
    print("\n================ EVALUATION BENCHMARK COMPARISON ================")
    print(f"{'Metric':<30} | {'Base Model':<15} | {'Fine-Tuned Model':<18} | {'Delta':<10}")
    print("-" * 80)
    for k in results["metrics"]["base"]:
        v_base = results["metrics"]["base"][k]
        v_ft = results["metrics"]["finetuned"][k]
        diff = round(v_ft - v_base, 3)
        diff_str = f"+{diff}" if diff > 0 else f"{diff}"
        print(f"{k:<30} | {v_base:<15} | {v_ft:<18} | {diff_str:<10}")
    print("=================================================================\n")

if __name__ == "__main__":
    evaluate_suite(
        data_dir="hf_bugbounty_suite/dataset/data",
        base_model_id="Qwen/Qwen2.5-0.5B-Instruct",
        ft_model_dir="hf_bugbounty_suite/model/bugbounty-triage-assistant-3b",
        num_samples=1
    )