import os
import numpy as np

MODEL_NAME = "distilbert-base-uncased"
OUTPUT_DIR = "./clickbait-lora-output"
DATASET_NAME = "marksverdhei/clickbait_title_classification"
MAX_LENGTH = 64
NUM_EPOCHS = 3
BATCH_SIZE = 16
LEARNING_RATE = 2e-4


# Tiny fallback dataset used ONLY if the HF dataset can't be downloaded
# (e.g. no internet in this sandbox). Replace with real data for your
# actual hackathon run - this is just so the script never crashes.
FALLBACK_TEXTS = [
    "You Won't Believe What This Celebrity Did Next",
    "Scientists Discover New Species of Frog in Amazon Rainforest",
    "This One Weird Trick Will Change Your Life Forever",
    "Local Council Approves Budget for New Public Library",
    "10 Shocking Secrets Doctors Don't Want You To Know",
    "Central Bank Raises Interest Rates by 0.25 Percent",
    "She Opened The Box And What Was Inside Left Everyone Speechless",
    "Annual Rainfall Report Shows Slight Increase Over Last Year",
    "Is Your Phone Secretly Spying On You? The Truth Will Shock You",
    "City Council Meeting Scheduled for Next Tuesday Evening",
    "This Simple Habit Could Add 10 Years To Your Life",
    "Quarterly GDP Figures Released By Statistics Department",
    "The Real Reason Celebrities Are Leaving Hollywood Will Stun You",
    "New Bridge Construction Project Begins in Downtown Area",
    "Doctors Hate Her! Find Out This Mom's Weight Loss Secret",
    "Ministry of Education Announces Updated Curriculum Guidelines",
    "You'll Never Guess Who Just Got Eliminated From The Show",
    "Regional Elections Results Certified By Election Commission",
    "This Photo Of A Dog Broke The Internet And Here's Why",
    "University Publishes Findings On Renewable Energy Efficiency",
]
FALLBACK_LABELS = [1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0]


def load_training_data():
    """Try to load the real dataset; fall back to tiny sample if offline."""
    try:
        from datasets import load_dataset

        ds = load_dataset(DATASET_NAME)
        train_split = ds["train"]

        # Normalize column names across dataset versions
        text_col = "title" if "title" in train_split.column_names else "text"
        label_col = "label" if "label" in train_split.column_names else "clickbait"

        texts = train_split[text_col]
        labels = train_split[label_col]

        print(f"Loaded {len(texts)} examples from '{DATASET_NAME}'")
        return list(texts), list(labels)

    except Exception as e:
        print(f"[WARN] Could not load '{DATASET_NAME}' ({e}).")
        print("[WARN] Falling back to tiny bundled sample dataset. "
              "Replace this with a real dataset before your demo!")
        return FALLBACK_TEXTS, FALLBACK_LABELS


def main():
    import torch
    from datasets import Dataset
    from transformers import (
        AutoTokenizer,
        AutoModelForSequenceClassification,
        TrainingArguments,
        Trainer,
        DataCollatorWithPadding,
    )
    from peft import LoraConfig, get_peft_model, TaskType

    texts, labels = load_training_data()

    dataset = Dataset.from_dict({"text": texts, "label": labels})
    dataset = dataset.train_test_split(test_size=0.15, seed=42)

    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

    def tokenize_fn(batch):
        return tokenizer(
            batch["text"],
            truncation=True,
            max_length=MAX_LENGTH,
        )

    tokenized = dataset.map(tokenize_fn, batched=True)

    base_model = AutoModelForSequenceClassification.from_pretrained(
        MODEL_NAME,
        num_labels=2,
        id2label={0: "not_clickbait", 1: "clickbait"},
        label2id={"not_clickbait": 0, "clickbait": 1},
    )

    # LoRA config - only trains small adapter matrices, base model frozen.
    # This is what makes it fast/cheap enough for a hackathon.
    lora_config = LoraConfig(
        task_type=TaskType.SEQ_CLS,
        r=8,
        lora_alpha=16,
        lora_dropout=0.1,
        target_modules=["q_lin", "v_lin"],  # DistilBERT attention projections
    )

    model = get_peft_model(base_model, lora_config)
    model.print_trainable_parameters()

    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

    def compute_metrics(eval_pred):
        logits, labels = eval_pred
        preds = np.argmax(logits, axis=-1)
        accuracy = (preds == labels).mean()
        return {"accuracy": accuracy}

    training_args = TrainingArguments(
        output_dir="./training-checkpoints",
        num_train_epochs=NUM_EPOCHS,
        per_device_train_batch_size=BATCH_SIZE,
        per_device_eval_batch_size=BATCH_SIZE,
        learning_rate=LEARNING_RATE,
        eval_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="accuracy",
        logging_steps=10,
        report_to=[],
    )

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

    trainer.train()

    eval_results = trainer.evaluate()
    print("\nFinal evaluation results:")
    print(eval_results)

    os.makedirs(OUTPUT_DIR, exist_ok=True)
    model.save_pretrained(OUTPUT_DIR)
    tokenizer.save_pretrained(OUTPUT_DIR)

    print(f"\nDone. LoRA adapter + tokenizer saved to: {OUTPUT_DIR}")
    print("Next step: copy this folder into backend/models/clickbait-lora/")
    print("in your TruthLens project, then restart the backend.")


if __name__ == "__main__":
    main()