import os
import pandas as pd
from huggingface_hub import login, whoami
from datasets import Dataset

from transformers import (
    AutoTokenizer, AutoModelForCausalLM,
    TrainingArguments, Trainer,
    DataCollatorForLanguageModeling
)

# ✅ 1. Load Hugging Face token securely
hf_token = os.environ.get("data3")
if not hf_token:
    raise ValueError("Missing HF token. Set 'data3' as a secret in your Space or Colab.")

# ✅ 2. Log in to Hugging Face
login(token=hf_token)

# ✅ 3. Get your username dynamically
username = whoami(token=hf_token)["name"]

# ✅ 4. Load your caption data
df = pd.read_csv("custom_captions.csv")

# ✅ 5. Format prompt
def to_prompt(row):
    return f"""Photo: {row['photo']}
Scenery: {row['scenery']}
Photo vibe: {row['photo_vibe']}
Vibe to convey: {row['desired_vibe']}
Captions user likes: "{row['liked_captions']}"
"""

# ✅ 6. Build full training text
df["text"] = df.apply(lambda row: to_prompt(row) + row["caption"], axis=1)

# ✅ 7. Convert to Hugging Face Dataset
dataset = Dataset.from_pandas(df[["text"]])

# ✅ 8. Load tokenizer and model (GPT-Neo)
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-125M")
tokenizer.pad_token = tokenizer.eos_token  # Required fix for GPT-style models
model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M")

# ✅ 9. Tokenize the dataset
def tokenize(example):
    return tokenizer(example["text"], truncation=True, padding="max_length", max_length=64)

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

# ✅ 10. Prepare data collator
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

# ✅ 11. Set training arguments
training_args = TrainingArguments(
    output_dir="./results",
    per_device_train_batch_size=2,
    num_train_epochs=1,
    save_strategy="epoch",
    push_to_hub=True,
    hub_model_id=f"{username}/caption-crafter-model",  # Dynamic + personalized
    hub_token=hf_token
)

# ✅ 12. Set up Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized,
    tokenizer=tokenizer,
    data_collator=data_collator
)

# ✅ 13. Train and push to Hugging Face Hub
trainer.train()
trainer.push_to_hub()