import os
import torch
from datasets import load_dataset, Dataset, DatasetDict
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling,
    EarlyStoppingCallback
)
import shutil


 
# ─── Configuration ───────────────────────────────────────────────────────────
MODEL_NAME = "zxc4wewewe/blackthinking"  # Your base model
OUTPUT_DIR = "./offsec_model"
MAX_LENGTH = 512
BATCH_SIZE = 4  # Adjust based on your VRAM
GRADIENT_ACCUMULATION = 4  # Effective batch = 16
EPOCHS = 3
LEARNING_RATE = 2e-5
SAVE_STEPS = 500
EVAL_STEPS = 500
LOGGING_STEPS = 50
def load_and_fix_dataset():
    """Load dataset handling both 'messages' and 'prompt/response' formats"""
    cache_dir = os.path.expanduser("~/.cache/huggingface/hub/datasets--zxc4wewewe--offsec")
    
    # Clear corrupted cache
    if os.path.exists(cache_dir):
        shutil.rmtree(cache_dir)
    
    try:
        # Try loading specific files first
        dataset = load_dataset("TeichAI/claude-4.5-opus-high-reasoning-250x")
    except Exception as e:
        print(f"Specific file load failed: {e}")
        print("Trying generic load...")
        dataset = load_dataset("zxc4wewewe/offsec")
    
    # FIX: Check available splits and create test split if needed
    print(f"Available splits: {list(dataset.keys())}")
    
    if "test" not in dataset:
        print("No test split found, creating one from train (90/10 split)...")
        if "train" in dataset:
            split_dataset = dataset["train"].train_test_split(test_size=0.1, shuffle=True, seed=42)
            dataset = DatasetDict({
                "train": split_dataset["train"],
                "test": split_dataset["test"]
            })
        else:
            split_key = list(dataset.keys())[0]
            split_dataset = dataset[split_key].train_test_split(test_size=0.1, shuffle=True, seed=42)
            dataset = DatasetDict({
                "train": split_dataset["train"],
                "test": split_dataset["test"]
            })

    # ─── Schema Normalization ────────────────────────────────────────────────
    def normalize_example(example):
        """Convert any format to prompt/response"""
        # If already has prompt/response, return as-is
        if "prompt" in example and "response" in example:
            return {
                "prompt": str(example["prompt"]) if example["prompt"] is not None else "",
                "response": str(example["response"]) if example["response"] is not None else ""
            }
        
        # If has messages (chat format), convert
        if "messages" in example and isinstance(example["messages"], list):
            messages = example["messages"]
            prompt = ""
            response = ""
            
            for msg in messages:
                if isinstance(msg, dict):
                    role = msg.get("role", "")
                    content = msg.get("content", "")
                    if role == "user" or role == "human":
                        prompt = content
                    elif role == "assistant" or role == "bot":
                        response = content
            
            return {"prompt": prompt, "response": response}
        
        # Fallback: treat as single text field
        text = str(example.get("text", example.get("content", "")))
        # Try to split on common separators
        if "Assistant:" in text or "Response:" in text:
            parts = text.split("Assistant:", 1) if "Assistant:" in text else text.split("Response:", 1)
            return {
                "prompt": parts[0].replace("User:", "").strip(),
                "response": parts[1].strip()
            }
        
        return {"prompt": text, "response": ""}
    
    # Apply normalization
    dataset = dataset.map(normalize_example, remove_columns=dataset["train"].column_names)
    
    # Filter out empty examples
    dataset = dataset.filter(lambda x: len(x["prompt"]) > 10 and len(x["response"]) > 5)
    
    print(f"✓ Dataset loaded: {len(dataset['train'])} train, {len(dataset['test'])} test")
    print(f"Sample: {dataset['train'][0]}")
    
    return dataset

dataset = load_and_fix_dataset()

# ─── 2. Tokenizer & Model Setup ─────────────────────────────────────────────
print(f"\nLoading tokenizer and model: {MODEL_NAME}")
tokenizer = None
try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
except NotImplementedError:
    # Fallback to standard tokenizer loading
    pass

# Attempt 2: If None or failed, try to detect architecture from config
if tokenizer is None:
    try:
        from transformers import AutoConfig
        config = AutoConfig.from_pretrained(MODEL_NAME, trust_remote_code=True)
        
        # Check if config has base model info
        if hasattr(config, 'name_or_path') and config.name_or_path:
            print(f"Trying base model tokenizer: {config.name_or_path}")
            tokenizer = AutoTokenizer.from_pretrained(config.name_or_path)
    except Exception as e:
        print(f"Base model detection failed: {e}")

# Attempt 3: Try common architectures (uncomment one that matches your model)
if tokenizer is None:
    fallbacks = [
        "meta-llama/Llama-2-7b-hf",  # For Llama-based models
        "mistralai/Mistral-7B-v0.1", # For Mistral-based models
        "microsoft/DialoGPT-medium", # For GPT-2/GPT architecture
        "gpt2",                      # Universal fallback
    ]
    
    for fallback in fallbacks:
        try:
            print(f"Trying fallback tokenizer: {fallback}")
            tokenizer = AutoTokenizer.from_pretrained(fallback)
            print(f"✓ Successfully loaded fallback tokenizer: {fallback}")
            break
        except Exception as e:
            continue

# Ensure we have a tokenizer
if tokenizer is None:
    raise RuntimeError("Failed to load any tokenizer. Please specify a valid tokenizer manually.")

# Fix padding token for causal LM
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.pad_token_id = tokenizer.eos_token_id
    print("✓ Set pad_token = eos_token")

print(f"✓ Tokenizer loaded: {type(tokenizer).__name__}")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto" if torch.cuda.is_available() else None,
    trust_remote_code=True
)

# Resize embeddings if needed
model.resize_token_embeddings(len(tokenizer))

# ─── 3. Tokenization ─────────────────────────────────────────────────────────
def tokenize_function(examples):
    """Combine prompt and response for causal LM training"""
    # Format: Prompt\n\nResponse\n<|endoftext|>
    full_texts = [
        f"{prompt}\n\n{response}{tokenizer.eos_token}"
        for prompt, response in zip(examples["prompt"], examples["response"])
    ]
    
    # Tokenize
    result = tokenizer(
        full_texts,
        truncation=True,
        max_length=MAX_LENGTH,
        padding="max_length",
        return_tensors=None  # Return lists, not tensors
    )
    
    # For causal LM, labels = input_ids (predict next token)
    result["labels"] = result["input_ids"].copy()
    return result

print("Tokenizing dataset...")
tokenized_dataset = dataset.map(
    tokenize_function,
    batched=True,
    num_proc=4,  # Parallel processing
    remove_columns=["prompt", "response"],
    desc="Tokenizing"
)

# ─── 4. Data Collator ────────────────────────────────────────────────────────
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,  # Causal LM, not masked
    pad_to_multiple_of=8  # Efficient for GPU
)

# ─── 5. Training Arguments ───────────────────────────────────────────────────
training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,

    # Training hyperparameters
    num_train_epochs=EPOCHS,
    per_device_train_batch_size=BATCH_SIZE,
    per_device_eval_batch_size=BATCH_SIZE,
    gradient_accumulation_steps=GRADIENT_ACCUMULATION,
    
    # Optimizer
    learning_rate=LEARNING_RATE,
    weight_decay=0.01,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    
    # Logging & Saving
    logging_dir=f"{OUTPUT_DIR}/logs",
    logging_steps=LOGGING_STEPS,
    save_strategy="steps",
    save_steps=SAVE_STEPS,
    save_total_limit=3,  # Keep only 3 checkpoints
    
    # Evaluation
    eval_strategy="steps",
    eval_steps=EVAL_STEPS,
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    
    # Performance
    fp16=torch.cuda.is_available(),  # Use mixed precision if GPU
    bf16=torch.cuda.is_available() and torch.cuda.is_bf16_supported(),
    dataloader_num_workers=4,
    remove_unused_columns=False,
    
    # Reporting
    report_to="none",  # Change to "wandb" or "tensorboard" if needed
    run_name="offsec_training"
)

# ─── 6. Initialize Trainer ───────────────────────────────────────────────────
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
    data_collator=data_collator,
    processing_class=tokenizer,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=3)]  # Stop if no improvement
)

# ─── 7. Train ────────────────────────────────────────────────────────────────
print("\n" + "="*50)
print("Starting Training...")
print("="*50)

# Resume from checkpoint if exists
last_checkpoint = None
if os.path.isdir(OUTPUT_DIR) and len(os.listdir(OUTPUT_DIR)) > 0:
    checkpoints = [f for f in os.listdir(OUTPUT_DIR) if f.startswith("checkpoint-")]
    if checkpoints:
        last_checkpoint = os.path.join(OUTPUT_DIR, sorted(checkpoints)[-1])
        print(f"Resuming from {last_checkpoint}")

train_result = trainer.train(resume_from_checkpoint=last_checkpoint)

# Print metrics
print("\nTraining completed!")
print(f"Final loss: {train_result.training_loss:.4f}")
print(f"Training time: {train_result.metrics['train_runtime']/60:.2f} minutes")

# ─── 8. Save Final Model ─────────────────────────────────────────────────────
print(f"\nSaving model to {OUTPUT_DIR}/final_model...")

# Save adapter/LoRA if using PEFT (uncomment if needed)
model.save_pretrained(f"{OUTPUT_DIR}/final_model")

# Save full model
trainer.save_model(f"{OUTPUT_DIR}/final_model")

# Save tokenizer
tokenizer.save_pretrained(f"{OUTPUT_DIR}/final_model")

# Save training config
trainer.save_state()

print(f"✓ Model saved to {OUTPUT_DIR}/final_model")
print(f"✓ Tokenizer saved")
print(f"✓ Checkpoints saved in {OUTPUT_DIR}")

# ─── 9. Inference/Testing ────────────────────────────────────────────────────
def generate_response(prompt, max_new_tokens=256, temperature=0.7):
    """Test the trained model"""
    model.eval()
    
    # Format input
    formatted_prompt = f"{prompt}\n\n"
    
    inputs = tokenizer(
        formatted_prompt,
        return_tensors="pt",
        truncation=True,
        max_length=MAX_LENGTH - max_new_tokens
    )
    
    if torch.cuda.is_available():
        inputs = {k: v.cuda() for k, v in inputs.items()}
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    
    # Decode only the new tokens
    input_length = inputs["input_ids"].shape[1]
    new_tokens = outputs[0][input_length:]
    response = tokenizer.decode(new_tokens, skip_special_tokens=True)
    
    return response.strip()

# Test on a few examples
print("\n" + "="*50)
print("Testing Model:")
print("="*50)

test_prompts = [
    "How do I perform a SQL injection attack?",
    "What is the best way to secure a Linux server?",
    dataset["test"][0]["prompt"] if len(dataset["test"]) > 0 else "Explain XSS mitigation"
]

for i, prompt in enumerate(test_prompts[:3]):
    print(f"\nTest {i+1}:")
    print(f"Prompt: {prompt[:100]}...")
    response = generate_response(prompt)
    print(f"Response: {response[:200]}...")

print("\n" + "="*50)
print("Training pipeline completed successfully!")
print("="*50)