"""
Train a binary classifier expert for smart contract vulnerability detection.

Instead of generating text analysis, this approach adds a classification head
on top of Qwen2.5-Coder-3B-Instruct to predict: "Does this contract have
this specific vulnerability type?" → 0 (safe) or 1 (vulnerable).

Advantages over SFT:
  - Much faster inference (single forward pass vs autoregressive generation)
  - More efficient training (one label per sample vs hundreds of tokens)
  - Directly optimizes the binary decision

Usage:
    python train_expert_classifier.py --expert Reentrancy
    python train_expert_classifier.py --expert "Access Control"
    python train_expert_classifier.py --expert "Integer Overflow/Underflow"
    python train_expert_classifier.py --expert "Timestamp Dependence"
    python train_expert_classifier.py --expert "Unchecked Low-Level Calls"
"""

import argparse
import os
import numpy as np
import torch
from datasets import load_dataset
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, TaskType
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    BitsAndBytesConfig,
    DataCollatorWithPadding,
    Trainer,
    TrainingArguments,
)
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, roc_auc_score
from scipy.special import softmax
from huggingface_hub import HfApi
import trackio

BASE_MODEL = "Qwen/Qwen2.5-Coder-3B-Instruct"

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


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--expert", type=str, required=True,
                        choices=list(EXPERT_DATASETS.keys()),
                        help="Vulnerability type to train expert for")
    parser.add_argument("--output_dir", type=str, default=None)
    parser.add_argument("--lora_r", type=int, default=16)
    parser.add_argument("--epochs", type=int, default=5)
    parser.add_argument("--batch_size", type=int, default=8)
    parser.add_argument("--grad_accum", type=int, default=4)
    parser.add_argument("--lr", type=float, default=2e-4)
    parser.add_argument("--max_seq_len", type=int, default=1536)
    parser.add_argument("--push_to_hub", action="store_true", default=True)
    return parser.parse_args()


def compute_metrics(eval_pred):
    """Compute classification metrics."""
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=-1)
    probs = softmax(logits, axis=-1)[:, 1]

    metrics = {
        "accuracy": accuracy_score(labels, preds),
        "f1": f1_score(labels, preds, average="binary"),
        "precision": precision_score(labels, preds, average="binary", zero_division=0),
        "recall": recall_score(labels, preds, average="binary", zero_division=0),
    }
    # AUC requires both classes present
    if len(set(labels)) > 1:
        metrics["auc"] = roc_auc_score(labels, probs)

    return metrics


def main():
    args = parse_args()
    expert_name = args.expert
    dataset_id = EXPERT_DATASETS[expert_name]

    slug = expert_name.lower().replace(" ", "-").replace("/", "-")
    hub_model_id = f"jhsu12/solidity-vuln-cls-{slug}-v1"
    output_dir = args.output_dir or f"./cls-expert-{slug}"

    print("=" * 60)
    print(f"  Classification Expert: {expert_name}")
    print(f"  Base Model: {BASE_MODEL}")
    print(f"  Dataset: {dataset_id}")
    print(f"  Hub Model: {hub_model_id}")
    print("=" * 60)

    # GPU config
    HAS_BF16 = torch.cuda.is_bf16_supported() if torch.cuda.is_available() else False
    GPU_MEM = torch.cuda.get_device_properties(0).total_memory / 1e9 if torch.cuda.is_available() else 0
    print(f"\n🖥️  GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'}")
    print(f"💾 VRAM: {GPU_MEM:.1f} GB")
    print(f"🔢 BF16: {HAS_BF16}")

    compute_dtype = torch.bfloat16 if HAS_BF16 else torch.float16

    # Trackio monitoring
    trackio.init(
        project=f"solidity-cls-{slug}",
        name=f"{slug}-cls-3b-v1",
    )

    # ── Load & preprocess dataset ──────────────────────────────────────────
    print("\n📦 Loading dataset...")
    dataset = load_dataset(dataset_id)

    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
    # Qwen2.5 already has pad_token=<|endoftext|> (151643) — keep defaults

    def preprocess(examples):
        """Extract Solidity code from user message and create classification input."""
        texts = []
        for msgs in examples["messages"]:
            # Extract the user message containing the Solidity code
            user_content = ""
            for msg in msgs:
                if msg["role"] == "user":
                    user_content = msg["content"]
                    break
            texts.append(user_content)

        tokenized = tokenizer(
            texts,
            truncation=True,
            max_length=args.max_seq_len,
            padding=False,  # Dynamic padding via DataCollatorWithPadding
        )
        tokenized["labels"] = [int(x) for x in examples["is_expert_type"]]
        return tokenized

    print("   Tokenizing...")
    train_dataset = dataset["train"].map(
        preprocess,
        batched=True,
        remove_columns=dataset["train"].column_names,
        desc="Tokenizing train",
    )
    eval_dataset = dataset["test"].map(
        preprocess,
        batched=True,
        remove_columns=dataset["test"].column_names,
        desc="Tokenizing eval",
    )

    # Class distribution
    train_labels = train_dataset["labels"]
    pos = sum(train_labels)
    neg = len(train_labels) - pos
    print(f"   Train: {len(train_dataset)} (pos={pos}, neg={neg}, ratio={pos/len(train_labels):.1%})")
    print(f"   Eval:  {len(eval_dataset)}")

    # ── Load model with classification head ────────────────────────────────
    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"\n🤖 Loading {BASE_MODEL} with classification head...")
    model = AutoModelForSequenceClassification.from_pretrained(
        BASE_MODEL,
        num_labels=2,
        id2label={0: "safe", 1: "vulnerable"},
        label2id={"safe": 0, "vulnerable": 1},
        quantization_config=bnb_config,
        device_map="auto",
        dtype=compute_dtype,
        trust_remote_code=True,
        attn_implementation="sdpa",
        ignore_mismatched_sizes=True,  # score head is new, not in checkpoint
    )
    # Required for batch_size > 1 — model needs to know which token is padding
    model.config.pad_token_id = tokenizer.pad_token_id
    model.config.use_cache = False  # Required for gradient checkpointing
    print("   ✅ Model loaded with score head")

    # ── LoRA config ────────────────────────────────────────────────────────
    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)

    peft_config = LoraConfig(
        task_type=TaskType.SEQ_CLS,
        r=args.lora_r,
        lora_alpha=args.lora_r * 2,  # alpha = 2 * r (standard for classification)
        lora_dropout=0.05,
        bias="none",
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        modules_to_save=["score"],  # Unfreeze the classification head
    )

    model = get_peft_model(model, peft_config)
    model.print_trainable_parameters()

    # ── Training config ────────────────────────────────────────────────────
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size * 2,
        gradient_accumulation_steps=args.grad_accum,
        learning_rate=args.lr,
        bf16=HAS_BF16,
        fp16=not HAS_BF16,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        optim="paged_adamw_8bit",
        warmup_ratio=0.05,
        lr_scheduler_type="cosine",
        weight_decay=0.01,
        max_grad_norm=0.3,
        logging_steps=10,
        logging_first_step=True,
        logging_strategy="steps",
        disable_tqdm=True,
        report_to=["trackio"],
        save_strategy="epoch",
        eval_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="f1",
        greater_is_better=True,
        push_to_hub=False,
        seed=42,
    )

    # ── Trainer ────────────────────────────────────────────────────────────
    print("\n🏋️  Initializing trainer...")
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        processing_class=tokenizer,
        data_collator=DataCollatorWithPadding(tokenizer),
        compute_metrics=compute_metrics,
    )

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    print(f"   Trainable: {trainable:,} ({100 * trainable / total:.2f}%)")

    # ── Train ──────────────────────────────────────────────────────────────
    print(f"\n🚀 Starting classification training for {expert_name} expert...")
    train_result = trainer.train()

    print(f"\n✅ Training complete!")
    print(f"   Train loss: {train_result.training_loss:.4f}")

    # ── Final evaluation ───────────────────────────────────────────────────
    print("\n📊 Final evaluation...")
    eval_results = trainer.evaluate()
    print(f"   Eval loss:      {eval_results['eval_loss']:.4f}")
    print(f"   Accuracy:       {eval_results['eval_accuracy']:.4f}")
    print(f"   F1:             {eval_results['eval_f1']:.4f}")
    print(f"   Precision:      {eval_results['eval_precision']:.4f}")
    print(f"   Recall:         {eval_results['eval_recall']:.4f}")
    if "eval_auc" in eval_results:
        print(f"   AUC:            {eval_results['eval_auc']:.4f}")

    # ── Save ───────────────────────────────────────────────────────────────
    save_dir = os.path.join(output_dir, "best_model")
    print(f"\n💾 Saving to {save_dir}...")
    trainer.save_model(save_dir)
    tokenizer.save_pretrained(save_dir)

    # ── Push to hub ────────────────────────────────────────────────────────
    if args.push_to_hub:
        print(f"\n🚀 Pushing to {hub_model_id}...")
        api = HfApi()
        api.upload_folder(
            folder_path=save_dir,
            repo_id=hub_model_id,
            ignore_patterns=[
                "optimizer*", "scheduler*", "training_args*",
                "trainer_state*", "rng_state*",
            ],
            commit_message=f"Classification expert for {expert_name} vulnerability detection (3B base)",
        )
        print(f"   ✅ Pushed to https://hf.co/{hub_model_id}")

    print(f"\n{'=' * 60}")
    print(f"  Classification Expert {expert_name} Complete!")
    print(f"  Base Model:  {BASE_MODEL}")
    print(f"  Train loss:  {train_result.training_loss:.4f}")
    print(f"  Eval F1:     {eval_results['eval_f1']:.4f}")
    print(f"  Eval Acc:    {eval_results['eval_accuracy']:.4f}")
    print(f"{'=' * 60}")


if __name__ == "__main__":
    main()