"""
Smart Contract Vulnerability Detector — Training Script
Works on H200, A100, L4, T4, or any CUDA GPU. Auto-configures.

🚀 Quick Start:
    pip install transformers trl peft datasets accelerate bitsandbytes
    huggingface-cli login
    python train_colab.py

flash-attn is OPTIONAL (speed boost only). If you have it, great. If not, no problem.
"""

import os
import torch
from datasets import load_dataset
from peft import LoraConfig
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from trl import SFTConfig, SFTTrainer

# ══════════════════════════════════════════════════════════════════════════════
# CONFIGURATION — Change these as needed
# ══════════════════════════════════════════════════════════════════════════════
MODEL_ID = "Qwen/Qwen2.5-Coder-7B-Instruct"
DATASET_ID = "jhsu12/solidity-vuln-detect-sft-data"
OUTPUT_DIR = "./solidity-vuln-detector"
HUB_MODEL_ID = "jhsu12/solidity-vulnerability-detector"  # Change to your username!

# ── Auto-detect GPU capabilities ─────────────────────────────────────────────
HAS_BF16 = torch.cuda.is_bf16_supported() if torch.cuda.is_available() else False
GPU_NAME = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU"
GPU_MEM = torch.cuda.get_device_properties(0).total_memory / 1e9 if torch.cuda.is_available() else 0

print(f"🖥️  GPU: {GPU_NAME}")
print(f"💾 VRAM: {GPU_MEM:.1f} GB")
print(f"🔢 BF16 support: {HAS_BF16}")

# Check if flash-attn is installed
HAS_FLASH_ATTN = False
try:
    import flash_attn
    HAS_FLASH_ATTN = True
    print(f"⚡ Flash Attention: v{flash_attn.__version__}")
except ImportError:
    print(f"⚡ Flash Attention: not installed (OK — using default attention, ~2x slower)")

# Auto-configure based on GPU memory
if GPU_MEM >= 120:  # H200 141GB / H100 80GB
    BATCH_SIZE = 8
    GRAD_ACCUM = 2
    LORA_R = 64
    MAX_SEQ_LEN = 2048
elif GPU_MEM >= 70:  # A100 80GB
    BATCH_SIZE = 4
    GRAD_ACCUM = 4
    LORA_R = 64
    MAX_SEQ_LEN = 2048
elif GPU_MEM >= 35:  # A100 40GB
    BATCH_SIZE = 2
    GRAD_ACCUM = 8
    LORA_R = 64
    MAX_SEQ_LEN = 2048
elif GPU_MEM >= 20:  # A10G 24GB, L4 24GB
    BATCH_SIZE = 1
    GRAD_ACCUM = 16
    LORA_R = 32
    MAX_SEQ_LEN = 1536
else:  # T4 16GB
    BATCH_SIZE = 1
    GRAD_ACCUM = 16
    LORA_R = 16
    MAX_SEQ_LEN = 1024

print(f"\n⚙️  Auto-config: batch={BATCH_SIZE}, grad_accum={GRAD_ACCUM}, "
      f"lora_r={LORA_R}, max_seq_len={MAX_SEQ_LEN}")
print(f"   Effective batch size: {BATCH_SIZE * GRAD_ACCUM}")

# ══════════════════════════════════════════════════════════════════════════════
# DATASET
# ══════════════════════════════════════════════════════════════════════════════
print("\n📦 Loading dataset...")
dataset = load_dataset(DATASET_ID)
train_dataset = dataset["train"]
eval_dataset = dataset["test"]
print(f"   Train: {len(train_dataset)} samples")
print(f"   Eval:  {len(eval_dataset)} samples")

# ══════════════════════════════════════════════════════════════════════════════
# MODEL
# ══════════════════════════════════════════════════════════════════════════════
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"\n🤖 Loading {MODEL_ID}...")
model_kwargs = dict(
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=compute_dtype,
    trust_remote_code=True,
)

# Use flash attention if available, otherwise sdpa (PyTorch built-in, still fast)
if HAS_FLASH_ATTN:
    model_kwargs["attn_implementation"] = "flash_attention_2"
    print("   Using: flash_attention_2")
else:
    model_kwargs["attn_implementation"] = "sdpa"
    print("   Using: sdpa (PyTorch native scaled dot-product attention)")

model = AutoModelForCausalLM.from_pretrained(MODEL_ID, **model_kwargs)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"
print(f"   ✅ Model loaded ({model.dtype})")

# ══════════════════════════════════════════════════════════════════════════════
# LORA CONFIG
# ══════════════════════════════════════════════════════════════════════════════
peft_config = LoraConfig(
    r=LORA_R,
    lora_alpha=16,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules="all-linear",
)

# ══════════════════════════════════════════════════════════════════════════════
# TRAINING CONFIG
# ══════════════════════════════════════════════════════════════════════════════
training_args = SFTConfig(
    output_dir=OUTPUT_DIR,
    
    # Training
    num_train_epochs=3,
    per_device_train_batch_size=BATCH_SIZE,
    gradient_accumulation_steps=GRAD_ACCUM,
    learning_rate=2e-4,
    
    # Precision
    bf16=HAS_BF16,
    fp16=not HAS_BF16,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
    
    # Sequence — "max_length" is the current TRL parameter name
    max_length=MAX_SEQ_LEN,
    packing=False,
    
    # Optimizer
    optim="paged_adamw_8bit",
    warmup_ratio=0.05,
    lr_scheduler_type="cosine",
    weight_decay=0.01,
    max_grad_norm=0.3,
    
    # Logging
    logging_steps=10,
    logging_first_step=True,
    logging_strategy="steps",
    report_to="none",
    
    # Saving
    save_strategy="epoch",
    eval_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    
    # Hub — automatically pushes LoRA adapters when training finishes
    push_to_hub=True,
    hub_model_id=HUB_MODEL_ID,
    hub_strategy="end",
    
    seed=42,
)

# ══════════════════════════════════════════════════════════════════════════════
# TRAIN
# ══════════════════════════════════════════════════════════════════════════════
print("\n🏋️ Initializing trainer...")
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    processing_class=tokenizer,
    peft_config=peft_config,
)

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}%)")
print(f"   Total:     {total:,}")

print("\n🚀 Starting training...")
train_result = trainer.train()

print(f"\n✅ Training complete!")
print(f"   Train loss: {train_result.training_loss:.4f}")
print(f"   Runtime: {train_result.metrics['train_runtime']/60:.1f} minutes")

# ── Evaluate ─────────────────────────────────────────────────────────────────
print("\n📊 Running evaluation...")
eval_results = trainer.evaluate()
print(f"   Eval loss: {eval_results['eval_loss']:.4f}")

# ── Save & Push ──────────────────────────────────────────────────────────────
print(f"\n💾 Saving & pushing to {HUB_MODEL_ID}...")
trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)
trainer.push_to_hub(commit_message="SFT+LoRA training: Solidity vulnerability detector")

print(f"\n{'='*60}")
print(f"🎉 Model available at: https://hf.co/{HUB_MODEL_ID}")
print(f"{'='*60}")

# ══════════════════════════════════════════════════════════════════════════════
# INFERENCE TESTS
# ══════════════════════════════════════════════════════════════════════════════
print("\n🔍 Running inference tests...")

test_code = '''pragma solidity ^0.8.0;

contract VulnerableBank {
    mapping(address => uint256) public balances;
    
    function deposit() public payable {
        balances[msg.sender] += msg.value;
    }
    
    function withdraw() public {
        uint256 bal = balances[msg.sender];
        require(bal > 0);
        (bool sent, ) = msg.sender.call{value: bal}("");
        require(sent, "Failed to send Ether");
        balances[msg.sender] = 0;
    }
}'''

messages = [
    {"role": "system", "content": "You are a smart contract security auditor. Analyze the provided Solidity code for vulnerabilities.\n\nFor each vulnerability found, provide:\n1. **Vulnerability Type**\n2. **Severity** — Critical, High, Medium, or Low\n3. **Location**\n4. **Explanation**\n5. **Recommendation**\n\nIf no vulnerabilities are found, state that the code appears safe."},
    {"role": "user", "content": f"Analyze the following Solidity code for vulnerabilities:\n\n```solidity\n{test_code}\n```"},
]

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

with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.1, do_sample=True, top_p=0.9)

response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
print("\n--- Test 1: Reentrancy Vulnerability (should detect!) ---")
print(response[:1500])
print("--- End ---")

# ── Test 2: Safe contract ────────────────────────────────────────────────────
safe_code = '''pragma solidity ^0.8.0;

import "@openzeppelin/contracts/security/ReentrancyGuard.sol";

contract SafeBank is ReentrancyGuard {
    mapping(address => uint256) public balances;
    
    function deposit() public payable {
        balances[msg.sender] += msg.value;
    }
    
    function withdraw() public nonReentrant {
        uint256 bal = balances[msg.sender];
        require(bal > 0, "No balance");
        balances[msg.sender] = 0;
        (bool sent, ) = msg.sender.call{value: bal}("");
        require(sent, "Failed to send Ether");
    }
}'''

messages2 = [
    {"role": "system", "content": messages[0]["content"]},
    {"role": "user", "content": f"Analyze the following Solidity code for vulnerabilities:\n\n```solidity\n{safe_code}\n```"},
]

text2 = tokenizer.apply_chat_template(messages2, tokenize=False, add_generation_prompt=True)
inputs2 = tokenizer(text2, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs2 = model.generate(**inputs2, max_new_tokens=512, temperature=0.1, do_sample=True, top_p=0.9)

response2 = tokenizer.decode(outputs2[0][inputs2["input_ids"].shape[-1]:], skip_special_tokens=True)
print("\n--- Test 2: Safe Contract (should say safe!) ---")
print(response2[:1500])
print("--- End ---")

print("\n✅ All done!")