import numpy as np
import torch
from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling,
)

# ─── Configuration ───────────────────────────────────────────────────────────
MODEL_NAME = "zxc4wewewe/blackthinking"                    # lightweight model suitable for CPU
MAX_LENGTH = 512                       # max token length per example
OUTPUT_DIR = "./results"
NUM_EPOCHS = 3
BATCH_SIZE = 2                         # small batch for CPU training
LEARNING_RATE = 5e-5
LOGGING_STEPS = 50

# ─── 1. Load dataset from Hugging Face Hub ───────────────────────────────────
dataset = load_dataset("zxc4wewewe/offsec")
print(f"Train: {len(dataset['train'])} examples | Test: {len(dataset['test'])} examples")
print(f"Columns: {dataset['train'].column_names}")


# ─── 2. Format & tokenize ────────────────────────────────────────────────────
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# GPT-2 has no pad token by default — use eos_token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token


def format_and_tokenize(examples):
    """Combine prompt + response into a single text and tokenize."""
    texts = [
        f"{prompt}{response}{tokenizer.eos_token}"
        for prompt, response in zip(examples["prompt"], examples["response"])
    ]
    tokenized = tokenizer(
        texts,
        truncation=True,
        max_length=MAX_LENGTH,
        padding="max_length",
    )
    # For causal LM, labels = input_ids (the model learns to predict next token)
    tokenized["labels"] = tokenized["input_ids"].copy()
    return tokenized


tokenized_dataset = dataset.map(
    format_and_tokenize,
    batched=True,
    remove_columns=dataset["train"].column_names,
    desc="Tokenizing",
)

print(f"Tokenized train: {len(tokenized_dataset['train'])} examples")


# ─── 3. Model ────────────────────────────────────────────────────────────────
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
model.resize_token_embeddings(len(tokenizer))

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,  # causal LM, not masked LM
)


# ─── 4. Training ─────────────────────────────────────────────────────────────
training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    overwrite_output_dir=True,
    num_train_epochs=NUM_EPOCHS,
    per_device_train_batch_size=BATCH_SIZE,
    per_device_eval_batch_size=BATCH_SIZE,
    eval_strategy="epoch",
    save_strategy="epoch",
    learning_rate=LEARNING_RATE,
    weight_decay=0.01,
    logging_dir="./logs",
    logging_steps=LOGGING_STEPS,
    load_best_model_at_end=True,
    save_total_limit=2,
    fp16=False,  # CPU-only
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
    data_collator=data_collator,
)

print("Starting training...")
trainer.train()

# Save final model
trainer.save_model(f"{OUTPUT_DIR}/final_model")
tokenizer.save_pretrained(f"{OUTPUT_DIR}/final_model")
print(f"Model saved to {OUTPUT_DIR}/final_model")


# ─── 5. Inference ────────────────────────────────────────────────────────────
def generate_response(prompt_text, max_new_tokens=256):
    """Generate a response given a prompt."""
    inputs = tokenizer(prompt_text, return_tensors="pt")
    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id,
        )
    # Decode only the generated part (skip the prompt tokens)
    generated = output_ids[0][inputs["input_ids"].shape[1]:]
    return tokenizer.decode(generated, skip_special_tokens=True)


# Example usage (uncomment to test after training):
sample_prompt = dataset["test"][0]["prompt"]
print("Prompt:", sample_prompt[:200], "...")
print("Generated:", generate_response(sample_prompt))