"""
Multi-Expert Evaluation Script for Smart Contract Vulnerability Detection.

Loads 5 expert LoRA adapters + 1 router (uses Integer Overflow as baseline router)
and classifies contracts by combining expert opinions.

Usage:
    python evaluate_experts.py --max_samples 200
    python evaluate_experts.py --max_new_tokens 384
"""

import argparse
import json
import os
import re
import time
import torch
from collections import defaultdict
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel

BASE_MODEL = "Qwen/Qwen2.5-Coder-7B-Instruct"
DATASET_ID = "jhsu12/solidity-vuln-detect-sft-data"

EXPERTS = {
    "Reentrancy": "jhsu12/solidity-vuln-expert-reentrancy-v1",
    "Access Control": "jhsu12/solidity-vuln-expert-access-control-v1",
    "Integer Overflow/Underflow": "jhsu12/solidity-vuln-expert-integer-overflow-underflow-v1",
    "Timestamp Dependence": "jhsu12/solidity-vuln-expert-timestamp-dependence-v1",
    "Unchecked Low-Level Calls": "jhsu12/solidity-vuln-expert-unchecked-low-level-calls-v1",
}

ALL_TYPES = list(EXPERTS.keys()) + ["tx.origin"]


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--max_samples", type=int, default=None)
    parser.add_argument("--max_new_tokens", type=int, default=256)
    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--output", type=str, default="expert_eval_results.json")
    parser.add_argument("--use_router", action="store_true", default=True,
                        help="Use routing: only query relevant experts based on type hints")
    return parser.parse_args()


def parse_expert_response(text):
    """Parse expert output: just need Vulnerable: Yes/No."""
    vuln_match = re.search(r'Vulnerable\s*[:\-]?\s*(Yes|No)', text, re.IGNORECASE)
    if vuln_match:
        return vuln_match.group(1).strip().lower() == "yes"
    # Fallback
    text_lower = text.lower()
    if "not" in text_lower and "vulnerable" in text_lower:
        return False
    if "yes" in text_lower and "vulnerable" in text_lower:
        return True
    return None


def parse_ground_truth(messages):
    """Extract vulnerability type from ground truth."""
    for msg in messages:
        if msg["role"] == "assistant":
            content = msg["content"]
            vuln_match = re.search(r'\*\*Vulnerable\*\*\s*[:\-]?\s*(Yes|No)', content, re.IGNORECASE)
            is_vuln = vuln_match.group(1).strip().lower() == "yes" if vuln_match else None
            type_match = re.search(r'\*\*Type\*\*\s*[:\-]?\s*(.+?)(?:\n|\r|$)', content)
            vtype = type_match.group(1).strip() if type_match else None
            sev_match = re.search(r'\*\*Severity\*\*\s*[:\-]?\s*(Critical|High|Medium|Low)', content, re.IGNORECASE)
            sev = sev_match.group(1).strip().capitalize() if sev_match else None
            return {"vulnerable": is_vuln, "type": vtype, "severity": sev}
    return None


def load_base_model():
    HAS_BF16 = torch.cuda.is_bf16_supported() if torch.cuda.is_available() else False
    compute_dtype = torch.bfloat16 if HAS_BF16 else torch.float16
    
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=compute_dtype,
        bnb_4bit_use_double_quant=True,
    )
    
    print(f"🤖 Loading base model {BASE_MODEL}...")
    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        quantization_config=bnb_config,
        device_map="auto",
        torch_dtype=compute_dtype,
        trust_remote_code=True,
        attn_implementation="sdpa",
    )
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"
    return model, tokenizer


def main():
    args = parse_args()
    
    print("=" * 60)
    print("  Multi-Expert Vulnerability Detection — Evaluation")
    print("=" * 60)
    
    # Load base + first expert for model loading
    model, tokenizer = load_base_model()
    
    # Load all expert adapters
    print(f"\n🔌 Loading expert adapters...")
    expert_models = {}
    for vtype, repo_id in EXPERTS.items():
        print(f"   Loading {vtype} from {repo_id}...")
        try:
            expert_models[vtype] = PeftModel.from_pretrained(model, repo_id)
            expert_models[vtype].eval()
            print(f"      ✅ Loaded")
        except Exception as e:
            print(f"      ❌ Failed: {e}")
    
    # Load dataset
    dataset = load_dataset(DATASET_ID, split="test")
    if args.max_samples:
        dataset = dataset.select(range(min(args.max_samples, len(dataset))))
    print(f"\n📦 Evaluating on {len(dataset)} samples")
    
    # Prepare prompts
    print(f"\n📝 Preparing prompts...")
    all_prompts = []
    all_truths = []
    for ex in dataset:
        gt = parse_ground_truth(ex["messages"])
        all_truths.append(gt)
        
        # Extract user prompt (system + user messages)
        prompt_msgs = [m for m in ex["messages"] if m["role"] != "assistant"]
        text = tokenizer.apply_chat_template(prompt_msgs, tokenize=False, add_generation_prompt=True)
        all_prompts.append(text)
    
    # Run inference per expert
    print(f"\n🔍 Running multi-expert inference...")
    all_predictions = []  # list of dicts: {expert_type: prediction}
    
    for expert_type, expert_model in expert_models.items():
        print(f"\n   [{expert_type}] Inference...")
        preds = []
        start = time.time()
        
        for i in range(0, len(all_prompts), args.batch_size):
            batch = all_prompts[i:i+args.batch_size]
            inputs = tokenizer(batch, return_tensors="pt", padding=True, truncation=True, max_length=1536).to(model.device)
            
            with torch.no_grad():
                outputs = expert_model.generate(
                    **inputs,
                    max_new_tokens=args.max_new_tokens,
                    do_sample=False,
                    pad_token_id=tokenizer.pad_token_id,
                )
            
            for j in range(len(batch)):
                input_len = inputs["attention_mask"][j].sum().item()
                response = tokenizer.decode(outputs[j][input_len:], skip_special_tokens=True)
                pred = parse_expert_response(response)
                preds.append(pred)
            
            if (i // args.batch_size + 1) % 5 == 0:
                elapsed = time.time() - start
                rate = (i + len(batch)) / elapsed
                print(f"      [{i+len(batch)}/{len(all_prompts)}] {rate:.1f} samples/s")
        
        print(f"      ✅ Done: {len(preds)} predictions")
        
        for idx, pred in enumerate(preds):
            if idx >= len(all_predictions):
                all_predictions.append({})
            all_predictions[idx][expert_type] = pred
    
    # Aggregate predictions
    print(f"\n🧠 Aggregating expert opinions...")
    final_predictions = []
    for idx, expert_votes in enumerate(all_predictions):
        # Simple voting: if any expert says Yes, classify as vulnerable with that type
        yes_experts = [et for et, pred in expert_votes.items() if pred == True]
        
        if yes_experts:
            # Pick the expert with highest confidence (or just first for now)
            final_pred = {
                "vulnerable": True,
                "vuln_type": yes_experts[0],
            }
        else:
            # No expert detected anything
            # Count None vs False
            no_experts = [et for et, pred in expert_votes.items() if pred == False]
            if no_experts:
                final_pred = {"vulnerable": False, "vuln_type": None}
            else:
                final_pred = {"vulnerable": None, "vuln_type": None}
        
        final_predictions.append(final_pred)
    
    # Compute metrics
    print(f"\n{'='*60}")
    print("  RESULTS")
    print(f"{'='*60}")
    
    # Binary metrics
    binary_preds = [p["vulnerable"] for p in final_predictions]
    binary_truths = [t["vulnerable"] if t else None for t in all_truths]
    
    valid_mask = [p is not None and t is not None for p, t in zip(binary_preds, binary_truths)]
    valid_p = [p for p, m in zip(binary_preds, valid_mask) if m]
    valid_t = [t for t, m in zip(binary_truths, valid_mask) if m]
    
    tp = sum(1 for p, t in zip(valid_p, valid_t) if p and t)
    tn = sum(1 for p, t in zip(valid_p, valid_t) if not p and not t)
    fp = sum(1 for p, t in zip(valid_p, valid_t) if p and not t)
    fn = sum(1 for p, t in zip(valid_p, valid_t) if not p and t)
    
    acc = (tp + tn) / (tp + tn + fp + fn) if (tp+tn+fp+fn) > 0 else 0
    prec = tp / (tp + fp) if (tp+fp) > 0 else 0
    rec = tp / (tp + fn) if (tp+fn) > 0 else 0
    f1 = 2 * prec * rec / (prec + rec) if (prec+rec) > 0 else 0
    
    print(f"\n📊 Binary Classification")
    print(f"   Accuracy:  {acc:.4f}")
    print(f"   Precision: {prec:.4f}")
    print(f"   Recall:    {rec:.4f}")
    print(f"   F1 Score:  {f1:.4f}")
    print(f"   TP={tp} TN={tn} FP={fp} FN={fn}")
    
    # Per-type metrics
    print(f"\n📊 Per Vulnerability Type")
    print(f"   {'Type':<30} {'Precision':>10} {'Recall':>10} {'F1':>10} {'Support':>8}")
    print(f"   {'-'*70}")
    
    for vtype in ALL_TYPES:
        type_preds = []
        type_truths = []
        for fpred, gt in zip(final_predictions, all_truths):
            if gt and gt["vulnerable"] == True and gt.get("type") == vtype:
                type_truths.append(vtype)
                if fpred.get("vuln_type") == vtype:
                    type_preds.append(vtype)
                else:
                    type_preds.append("Other/None")
        
        if type_truths:
            tp_type = sum(1 for p, t in zip(type_preds, type_truths) if p == vtype)
            fp_type = sum(1 for p in type_preds if p == vtype) - tp_type
            fn_type = len(type_truths) - tp_type
            
            p = tp_type / (tp_type + fp_type) if (tp_type + fp_type) > 0 else 0
            r = tp_type / (tp_type + fn_type) if (tp_type + fn_type) > 0 else 0
            f1_type = 2 * p * r / (p + r) if (p + r) > 0 else 0
            
            print(f"   {vtype:<30} {p:>10.4f} {r:>10.4f} {f1_type:>10.4f} {len(type_truths):>8}")
    
    # Save results
    results = {
        "num_samples": len(dataset),
        "binary_metrics": {
            "accuracy": round(acc, 4),
            "precision": round(prec, 4),
            "recall": round(rec, 4),
            "f1": round(f1, 4),
            "tp": tp, "tn": tn, "fp": fp, "fn": fn,
        },
        "expert_predictions": [{"predictions": p, "truth": {"vulnerable": t["vulnerable"] if t else None, "type": t["type"] if t else None}}
                            for p, t in zip(all_predictions, all_truths)],
    }
    
    with open(args.output, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\n💾 Results saved to {args.output}")


if __name__ == "__main__":
    main()