import torch
from transformers import (
    AutoTokenizer, 
    AutoModelForSequenceClassification,
    Trainer, 
    TrainingArguments
)
from datasets import load_from_disk
import numpy as np
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
import os

print("="*60)
print("RETRAINING WITH ENHANCED DATASET")
print("="*60)

# Load enhanced dataset
print("📂 Loading enhanced dataset...")
dataset = load_from_disk("enhanced_data/hf_dataset")

print(f"Train size: {len(dataset['train'])}")
print(f"Test size: {len(dataset['test'])}")

# Load tokenizer and model
print("🤖 Loading CodeBERT...")
model_name = "microsoft/codebert-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(
    model_name, 
    num_labels=2
)

# Tokenization function
def tokenize_function(examples):
    return tokenizer(
        examples["code"],
        padding="max_length",
        truncation=True,
        max_length=256  # Increased for better context
    )

print("🔢 Tokenizing dataset...")
tokenized_datasets = dataset.map(tokenize_function, batched=True)

# Remove text columns
columns_to_remove = ["code", "type", "explanation", "has_syntax_error", "syntax_error"]
columns_to_remove = [col for col in columns_to_remove if col in tokenized_datasets["train"].column_names]
tokenized_datasets = tokenized_datasets.remove_columns(columns_to_remove)
tokenized_datasets.set_format("torch")

# Training arguments (better this time)
training_args = TrainingArguments(
    output_dir="./enhanced_model",
    num_train_epochs=5,  # More epochs
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir="./enhanced_logs",
    logging_steps=50,
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    greater_is_better=True,
    save_total_limit=2,
    report_to="none"
)

# Better metrics
def compute_metrics(p):
    predictions, labels = p
    predictions = np.argmax(predictions, axis=1)
    
    accuracy = accuracy_score(labels, predictions)
    precision, recall, f1, _ = precision_recall_fscore_support(
        labels, predictions, average="binary", zero_division=0
    )
    
    return {
        "accuracy": accuracy,
        "precision": precision,
        "recall": recall,
        "f1": f1
    }

# Create trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["test"],
    compute_metrics=compute_metrics,
)

# Train
print("🚀 Training enhanced model...")
print("This will take 10-15 minutes...")
trainer.train()

# Evaluate
print("\n📊 Final Evaluation:")
metrics = trainer.evaluate()
print(f"Accuracy: {metrics['eval_accuracy']:.2%}")
print(f"Precision: {metrics['eval_precision']:.2%}")
print(f"Recall: {metrics['eval_recall']:.2%}")
print(f"F1 Score: {metrics['eval_f1']:.2%}")

# Save model
print("\n💾 Saving enhanced model...")
trainer.save_model("enhanced_saved_model")
tokenizer.save_pretrained("enhanced_saved_model")

print("\n" + "="*60)
print("🎉 ENHANCED MODEL TRAINED!")
print("Model saved to: enhanced_saved_model/")
print("="*60)

# Quick test
print("\n🔍 Quick test:")
test_codes = [
    """query = f"SELECT * FROM users WHERE id = {user_id}" """,
    """api_key = os.getenv("API_KEY")""",
    """def test()\n    print("hello")""",  # Syntax error
]

for code in test_codes:
    inputs = tokenizer(code, return_tensors="pt", truncation=True, max_length=256)
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    
    prediction = "VULNERABLE" if probs[0][1] > 0.5 else "SAFE"
    print(f"\nCode: {code[:50]}...")
    print(f"  Prediction: {prediction}")
    print(f"  Confidence: Safe={probs[0][0]:.2%}, Vulnerable={probs[0][1]:.2%}")