"""
Baseline Evaluation: Qwen2.5-Coder-3B-Instruct on Solidity Vulnerability Detection

Evaluates the model's zero-shot ability to detect a specific vulnerability type
(e.g., reentrancy) in Solidity smart contracts. Designed to be reusable across
different vulnerability datasets by changing --dataset_id and --vulnerability_name.

Key features:
  - Binary True/False output with retry on parse failure (up to --max_retries)
  - Per-sample inference time tracking
  - Saves summary JSON + detailed per-sample JSONL for error analysis

Usage:
    python evaluate_baseline.py \
        --model_id Qwen/Qwen2.5-Coder-3B-Instruct \
        --dataset_id jhsu12/solidity-vuln-expert-reentrancy \
        --vulnerability_name reentrancy \
        --split test \
        --max_retries 3
"""

import argparse
import json
import re
import time
from pathlib import Path

import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from sklearn.metrics import (
    accuracy_score,
    precision_score,
    recall_score,
    f1_score,
    confusion_matrix,
    classification_report,
)


# ---------------------------------------------------------------------------
# 1. Argument parsing
# ---------------------------------------------------------------------------
def parse_args():
    p = argparse.ArgumentParser(description="Baseline eval for solidity vulnerability detection")
    p.add_argument("--model_id", type=str, default="Qwen/Qwen2.5-Coder-3B-Instruct")
    p.add_argument("--dataset_id", type=str, default="jhsu12/solidity-vuln-expert-reentrancy")
    p.add_argument("--vulnerability_name", type=str, default="reentrancy",
                    help="Name of the vulnerability type (used in the prompt)")
    p.add_argument("--split", type=str, default="test",
                    help="Dataset split to evaluate on")
    p.add_argument("--max_samples", type=int, default=None,
                    help="Limit the number of samples (None = use full split)")
    p.add_argument("--max_new_tokens", type=int, default=16,
                    help="Max tokens to generate per attempt")
    p.add_argument("--max_retries", type=int, default=3,
                    help="Max retry attempts when model doesn't output True/False")
    p.add_argument("--output_dir", type=str, default="./eval_results",
                    help="Directory to save results")
    p.add_argument("--device", type=str, default=None,
                    help="Force device (auto-detected if omitted)")
    return p.parse_args()


# ---------------------------------------------------------------------------
# 2. Extract Solidity code from the dataset's user message
# ---------------------------------------------------------------------------
def extract_solidity_code(messages: list[dict]) -> str:
    """Pull the user-role content from the chat messages list."""
    for msg in messages:
        if msg["role"] == "user":
            return msg["content"]
    return ""


# ---------------------------------------------------------------------------
# 3. Build prompts (initial + retry variants)
# ---------------------------------------------------------------------------
def build_prompt(solidity_code: str, vulnerability_name: str, attempt: int = 0) -> list[dict]:
    """
    Construct a chat prompt asking for a binary True/False answer.

    attempt=0: Standard prompt
    attempt=1: Adds explicit reminder about output format
    attempt=2+: Even more constrained, tells model its prior answer was invalid
    """
    system_msg = (
        f"You are a smart contract security expert specializing in {vulnerability_name} vulnerabilities. "
        f"Your task is to determine whether the given Solidity code contains a {vulnerability_name} vulnerability. "
        "You must respond with ONLY one word: 'True' if the code is vulnerable, or 'False' if it is not. "
        "Do not provide any explanation, analysis, or additional text. Just output 'True' or 'False'."
    )

    user_msg = (
        f"Analyze the following Solidity code for {vulnerability_name} vulnerability:\n\n"
        f"{solidity_code}"
    )

    messages = [
        {"role": "system", "content": system_msg},
        {"role": "user", "content": user_msg},
    ]

    if attempt == 1:
        # Add a stronger format reminder as an assistant prefill nudge
        messages.append({"role": "assistant", "content": "After analyzing the code, my answer is:"})
        messages.append({"role": "user", "content": "Remember: respond with ONLY 'True' or 'False'. Nothing else."})
    elif attempt >= 2:
        # Simulate a correction conversation
        messages.append({"role": "assistant", "content": "After analyzing the code, my answer is:"})
        messages.append({
            "role": "user",
            "content": (
                "Your previous response was not in the correct format. "
                "You MUST reply with exactly one word: 'True' or 'False'. "
                "Is there a " + vulnerability_name + " vulnerability? Reply now:"
            ),
        })

    return messages


# ---------------------------------------------------------------------------
# 4. Parse model output into a boolean prediction
# ---------------------------------------------------------------------------
def parse_prediction(text: str) -> bool | None:
    """
    Robustly parse model output to True/False.
    Returns None if the output is unparseable.
    """
    cleaned = text.strip().lower()
    if not cleaned:
        return None

    # Direct match
    if cleaned in ("true", "true.", "true,"):
        return True
    if cleaned in ("false", "false.", "false,"):
        return False

    # Check if the first word is true/false
    first_word = re.split(r"[\s.,;:!?\n]", cleaned)[0]
    if first_word == "true":
        return True
    if first_word == "false":
        return False

    # Fallback: search for true/false anywhere (take the first occurrence)
    true_pos = cleaned.find("true")
    false_pos = cleaned.find("false")

    if true_pos != -1 and (false_pos == -1 or true_pos < false_pos):
        return True
    if false_pos != -1 and (true_pos == -1 or false_pos < true_pos):
        return False

    # Heuristics: "yes"/"vulnerable" → True, "no"/"not vulnerable" → False
    if any(kw in cleaned for kw in ["not vulnerable", "no vulnerability", "is not"]):
        return False
    if any(kw in cleaned for kw in ["vulnerable", "yes"]):
        return True

    return None  # Unparseable


# ---------------------------------------------------------------------------
# 5. Single-sample inference with retry
# ---------------------------------------------------------------------------
def infer_single(
    model,
    tokenizer,
    solidity_code: str,
    vulnerability_name: str,
    device: torch.device,
    max_new_tokens: int,
    max_retries: int,
) -> dict:
    """
    Run inference on a single sample. If the model's output can't be parsed
    as True/False, retry up to max_retries times with increasingly explicit
    prompts. Returns a dict with prediction, raw outputs, timing, and retry count.
    """
    all_raw_outputs = []       # every attempt's raw text
    prediction = None
    total_inference_time = 0.0

    for attempt in range(max_retries + 1):  # attempt 0 = first try, then up to max_retries
        messages = build_prompt(solidity_code, vulnerability_name, attempt=attempt)

        text = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        inputs = tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            max_length=4096,
        ).to(device)

        # --- Timed generation ---
        t0 = time.perf_counter()
        with torch.no_grad():
            output_ids = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                temperature=None,
                top_p=None,
                pad_token_id=tokenizer.pad_token_id,
            )
        # Sync GPU clock before stopping timer
        if device.type == "cuda":
            torch.cuda.synchronize()
        t1 = time.perf_counter()

        inference_time = t1 - t0
        total_inference_time += inference_time

        # Decode only newly generated tokens
        input_len = inputs["input_ids"].shape[1]
        generated_ids = output_ids[0][input_len:]
        raw_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
        all_raw_outputs.append(raw_text)

        prediction = parse_prediction(raw_text)
        if prediction is not None:
            break  # Successfully parsed → stop retrying

    return {
        "prediction": prediction,                 # bool or None if all retries failed
        "raw_outputs": all_raw_outputs,            # list of strings, one per attempt
        "attempts": len(all_raw_outputs),          # 1 = first try succeeded
        "inference_time": round(total_inference_time, 4),  # seconds (sum of all attempts)
    }


# ---------------------------------------------------------------------------
# 6. Main evaluation loop
# ---------------------------------------------------------------------------
def main():
    args = parse_args()

    # --- Device setup ---
    if args.device:
        device = torch.device(args.device)
    elif torch.cuda.is_available():
        device = torch.device("cuda")
    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        device = torch.device("mps")
    else:
        device = torch.device("cpu")
    print(f"Device: {device}")

    # --- Load model & tokenizer ---
    print(f"Loading model: {args.model_id}")
    tokenizer = AutoTokenizer.from_pretrained(args.model_id, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    dtype = torch.float16 if device.type == "cuda" else torch.float32
    model = AutoModelForCausalLM.from_pretrained(
        args.model_id,
        torch_dtype=dtype,
        device_map=device.type if device.type == "cuda" else None,
        trust_remote_code=True,
    )
    if device.type != "cuda":
        model = model.to(device)
    model.eval()
    print("Model loaded.\n")

    # --- Load dataset ---
    print(f"Loading dataset: {args.dataset_id} (split={args.split})")
    ds = load_dataset(args.dataset_id, split=args.split)
    if args.max_samples is not None:
        ds = ds.select(range(min(args.max_samples, len(ds))))
    total = len(ds)
    print(f"Evaluating on {total} samples (max_retries={args.max_retries})\n")

    # --- Per-sample inference ---
    ground_truths = []       # bool
    predictions = []         # bool | None
    sample_results = []      # full per-sample dicts
    parse_failures = 0
    total_retries = 0

    wall_start = time.time()

    for idx in range(total):
        row = ds[idx]
        gt = bool(row["is_expert_type"])
        ground_truths.append(gt)

        code = extract_solidity_code(row["messages"])

        result = infer_single(
            model=model,
            tokenizer=tokenizer,
            solidity_code=code,
            vulnerability_name=args.vulnerability_name,
            device=device,
            max_new_tokens=args.max_new_tokens,
            max_retries=args.max_retries,
        )

        pred = result["prediction"]
        predictions.append(pred)

        if pred is None:
            parse_failures += 1
        if result["attempts"] > 1:
            total_retries += result["attempts"] - 1

        sample_results.append({
            "index": idx,
            "ground_truth": gt,
            "prediction": pred,
            "raw_outputs": result["raw_outputs"],
            "attempts": result["attempts"],
            "inference_time": result["inference_time"],
        })

        # Progress every 20 samples or at the end
        done = idx + 1
        if done % 20 == 0 or done == total:
            elapsed = time.time() - wall_start
            avg_time = elapsed / done
            eta = avg_time * (total - done)
            print(
                f"  [{done:>4}/{total}]  "
                f"elapsed={elapsed:.1f}s  "
                f"avg={avg_time:.2f}s/sample  "
                f"ETA={eta:.0f}s  "
                f"parse_fails={parse_failures}  "
                f"retries={total_retries}"
            )

    wall_time = time.time() - wall_start

    # ---------------------------------------------------------------------------
    # 7. Resolve unparseable predictions
    # ---------------------------------------------------------------------------
    # After all retries exhausted, if still None → assign opposite of ground truth
    # (guarantees it's counted as wrong). Reported separately so you know.
    preds_resolved = []
    for pred, gt in zip(predictions, ground_truths):
        if pred is None:
            preds_resolved.append(not gt)
        else:
            preds_resolved.append(pred)

    # ---------------------------------------------------------------------------
    # 8. Compute metrics
    # ---------------------------------------------------------------------------
    acc = accuracy_score(ground_truths, preds_resolved)
    prec = precision_score(ground_truths, preds_resolved, zero_division=0)
    rec = recall_score(ground_truths, preds_resolved, zero_division=0)
    f1 = f1_score(ground_truths, preds_resolved, zero_division=0)
    cm = confusion_matrix(ground_truths, preds_resolved, labels=[False, True])
    report = classification_report(
        ground_truths, preds_resolved,
        target_names=["Not Vulnerable", "Vulnerable"],
        zero_division=0,
    )

    # Inference time stats
    sample_times = [r["inference_time"] for r in sample_results]
    avg_time = sum(sample_times) / len(sample_times) if sample_times else 0
    min_time = min(sample_times) if sample_times else 0
    max_time = max(sample_times) if sample_times else 0
    median_time = sorted(sample_times)[len(sample_times) // 2] if sample_times else 0

    # ---------------------------------------------------------------------------
    # 9. Print results
    # ---------------------------------------------------------------------------
    print("\n" + "=" * 64)
    print("  BASELINE EVALUATION RESULTS")
    print("=" * 64)
    print(f"  Model:           {args.model_id}")
    print(f"  Dataset:         {args.dataset_id}")
    print(f"  Vulnerability:   {args.vulnerability_name}")
    print(f"  Split:           {args.split} ({total} samples)")
    print(f"  Max retries:     {args.max_retries}")
    print("-" * 64)
    print(f"  Accuracy:        {acc:.4f}")
    print(f"  Precision:       {prec:.4f}  (of predicted vuln, how many truly are)")
    print(f"  Recall:          {rec:.4f}  (of truly vuln, how many we detected)")
    print(f"  F1 Score:        {f1:.4f}")
    print("-" * 64)
    print(f"  Parse failures:  {parse_failures}/{total} "
          f"({parse_failures/total*100:.1f}% — still unparseable after all retries)")
    print(f"  Total retries:   {total_retries} (across all samples)")
    print("-" * 64)
    print(f"  Inference time (per sample, incl. retries):")
    print(f"    Mean:   {avg_time:.3f}s")
    print(f"    Median: {median_time:.3f}s")
    print(f"    Min:    {min_time:.3f}s")
    print(f"    Max:    {max_time:.3f}s")
    print(f"  Total wall time: {wall_time:.1f}s")
    print("-" * 64)
    print(f"  Confusion Matrix (rows=actual, cols=predicted):")
    print(f"                      Pred:Not Vuln  Pred:Vuln")
    print(f"  Actual:Not Vuln     {cm[0][0]:>12}  {cm[0][1]:>9}")
    print(f"  Actual:Vuln         {cm[1][0]:>12}  {cm[1][1]:>9}")
    print(f"\n{report}")

    # ---------------------------------------------------------------------------
    # 10. Save results
    # ---------------------------------------------------------------------------
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # --- Summary JSON ---
    summary = {
        "model_id": args.model_id,
        "dataset_id": args.dataset_id,
        "vulnerability_name": args.vulnerability_name,
        "split": args.split,
        "num_samples": total,
        "max_retries": args.max_retries,
        "metrics": {
            "accuracy": round(acc, 4),
            "precision": round(prec, 4),
            "recall": round(rec, 4),
            "f1_score": round(f1, 4),
        },
        "confusion_matrix": {
            "true_negative": int(cm[0][0]),
            "false_positive": int(cm[0][1]),
            "false_negative": int(cm[1][0]),
            "true_positive": int(cm[1][1]),
        },
        "label_distribution": {
            "vulnerable": sum(ground_truths),
            "not_vulnerable": total - sum(ground_truths),
        },
        "parse_failures": parse_failures,
        "parse_failure_rate": round(parse_failures / total, 4),
        "total_retries": total_retries,
        "inference_time": {
            "wall_time_seconds": round(wall_time, 2),
            "per_sample_mean": round(avg_time, 4),
            "per_sample_median": round(median_time, 4),
            "per_sample_min": round(min_time, 4),
            "per_sample_max": round(max_time, 4),
        },
    }
    summary_path = output_dir / f"baseline_{args.vulnerability_name}_summary.json"
    with open(summary_path, "w") as f:
        json.dump(summary, f, indent=2)
    print(f"\nSummary saved to:  {summary_path}")

    # --- Detailed per-sample JSONL ---
    detail_path = output_dir / f"baseline_{args.vulnerability_name}_detailed.jsonl"
    with open(detail_path, "w") as f:
        for r in sample_results:
            r_out = {
                "index": r["index"],
                "ground_truth": r["ground_truth"],
                "prediction": r["prediction"],
                "prediction_resolved": preds_resolved[r["index"]],
                "correct": r["ground_truth"] == preds_resolved[r["index"]],
                "attempts": r["attempts"],
                "inference_time": r["inference_time"],
                "raw_outputs": r["raw_outputs"],
            }
            f.write(json.dumps(r_out) + "\n")
    print(f"Details saved to:  {detail_path}")

    # --- Show sample parse failures ---
    if parse_failures > 0:
        print(f"\nSample parse failures (up to 5):")
        shown = 0
        for r in sample_results:
            if r["prediction"] is None and shown < 5:
                print(f"  [idx={r['index']}] attempts={r['attempts']}  "
                      f"raw_outputs={r['raw_outputs']}")
                shown += 1


if __name__ == "__main__":
    main()