"""
Smart Contract Vulnerability Detector — V2 Training Script
Improvements over v1:
  - Dataset v2: removed 1,512 generic no-type samples, added 719 samples from darkknight25
  - Better type balance: Access Control 61→219, Unchecked Calls 67→215, Reentrancy 616→767
  - Explicit model save + push to avoid the v1 push_to_hub silent failure

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_v2.py
"""

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

# ══════════════════════════════════════════════════════════════════════════════
# CONFIGURATION
# ══════════════════════════════════════════════════════════════════════════════
MODEL_ID = "Qwen/Qwen2.5-Coder-7B-Instruct"
DATASET_ID = "jhsu12/solidity-vuln-detect-sft-data"
OUTPUT_DIR = "./solidity-vuln-detector-v2"
HUB_MODEL_ID = "jhsu12/solidity-vulnerability-detector"

# ── 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}")

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 (using sdpa)")

# Auto-configure based on GPU memory
if GPU_MEM >= 120:  # H200/H100
    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/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 (v2)...")
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,
)

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")

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=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",
    disable_tqdm=True,
    report_to="none",
    
    # Saving — save every epoch + keep best
    save_strategy="epoch",
    eval_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    
    # Hub — we'll push manually to avoid silent failures
    push_to_hub=False,
    
    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 best model locally ──────────────────────────────────────────────────
SAVE_DIR = os.path.join(OUTPUT_DIR, "best_model")
print(f"\n💾 Saving best model to {SAVE_DIR}...")
trainer.save_model(SAVE_DIR)
tokenizer.save_pretrained(SAVE_DIR)

# Verify files exist
import glob
saved_files = glob.glob(os.path.join(SAVE_DIR, "*"))
print(f"   Saved {len(saved_files)} files:")
for f in sorted(saved_files):
    size = os.path.getsize(f) / 1e6
    print(f"     {os.path.basename(f)} ({size:.1f} MB)")

# ── Push to Hub explicitly ────────────────────────────────────────────────────
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="v2: Improved dataset — removed generic samples, added Access Control/Reentrancy/Unchecked Calls from darkknight25",
)
print(f"   ✅ Pushed to https://hf.co/{HUB_MODEL_ID}")

# ── Print training log summary ────────────────────────────────────────────────
print(f"\n{'='*60}")
print(f"  V2 TRAINING SUMMARY")
print(f"{'='*60}")

# Get eval losses from trainer state
import json
state_path = os.path.join(OUTPUT_DIR, "trainer_state.json")
if os.path.exists(state_path):
    # Try each checkpoint's state
    pass

for entry in trainer.state.log_history:
    if "eval_loss" in entry:
        print(f"  Epoch {entry.get('epoch', '?')}: eval_loss={entry['eval_loss']:.4f}")

print(f"\n  Best model: {trainer.state.best_model_checkpoint}")
print(f"  Best eval loss: {trainer.state.best_metric:.4f}")
print(f"  Train loss: {train_result.training_loss:.4f}")
print(f"{'='*60}")

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

SYSTEM = ("You are a smart contract security auditor. Analyze the provided Solidity code for vulnerabilities.\n\n"
          "For each vulnerability found, provide:\n1. **Vulnerability Type**\n2. **Severity** — Critical, High, Medium, or Low\n"
          "3. **Location**\n4. **Explanation**\n5. **Recommendation**\n\nIf no vulnerabilities are found, state that the code appears safe.")

# Test 1: Reentrancy
test1 = '''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;
    }
}'''

# Test 2: Access Control (v1 weakness!)
test2 = '''pragma solidity ^0.8.0;
contract VulnerableToken {
    address public owner;
    uint256 public totalSupply;
    mapping(address => uint256) public balances;
    constructor() { owner = msg.sender; }
    function mint(address to, uint256 amount) public {
        totalSupply += amount;
        balances[to] += amount;
    }
    function setOwner(address newOwner) public {
        owner = newOwner;
    }
}'''

# Test 3: Safe contract
test3 = '''pragma solidity ^0.8.0;
import "@openzeppelin/contracts/security/ReentrancyGuard.sol";
import "@openzeppelin/contracts/access/Ownable.sol";
contract SafeBank is ReentrancyGuard, Ownable {
    mapping(address => uint256) public balances;
    constructor() Ownable(msg.sender) {}
    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");
    }
}'''

tests = [
    ("Reentrancy (should detect!)", test1),
    ("Access Control (should detect! — v1 weakness)", test2),
    ("Safe Contract (should say safe!)", test3),
]

for name, code in tests:
    messages = [
        {"role": "system", "content": SYSTEM},
        {"role": "user", "content": f"Analyze the following Solidity code for vulnerabilities:\n\n```solidity\n{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(f"\n--- {name} ---")
    print(response[:1000])
    print("--- End ---")

print("\n✅ All done! Model pushed to https://hf.co/" + HUB_MODEL_ID)