import os
import torch
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
)
from peft import LoraConfig
from trl import SFTTrainer
from huggingface_hub import HfFolder

# -----------------------------
# Environment & Memory Tweaks
# -----------------------------
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
os.environ["HF_HOME"] = "/tmp/huggingface"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface/transformers"
os.environ["HF_DATASETS_CACHE"] = "/tmp/huggingface/datasets"
os.environ["HF_METRICS_CACHE"] = "/tmp/huggingface/metrics"

os.makedirs("/tmp/huggingface", exist_ok=True)

# -----------------------------
# Model & Dataset Configuration
# -----------------------------
model_name = "meta-llama/Llama-3.2-3B-Instruct"
dataset_name = "Krishnasri2027/guanaco-llama2-2k"
output_dir = "/tmp/results"

# -----------------------------
# LoRA Parameters
# -----------------------------
lora_r = 64
lora_alpha = 16
lora_dropout = 0.1

# -----------------------------
# BitsAndBytes Configuration
# -----------------------------
use_4bit = True
bnb_4bit_compute_dtype = "float16"
bnb_4bit_quant_type = "nf4"
use_nested_quant = False

compute_dtype = getattr(torch, bnb_4bit_compute_dtype)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=use_4bit,
    bnb_4bit_quant_type=bnb_4bit_quant_type,
    bnb_4bit_compute_dtype=compute_dtype,
    bnb_4bit_use_double_quant=use_nested_quant,
)

# -----------------------------
# Dataset Loading
# -----------------------------
dataset = load_dataset(dataset_name, split="train", cache_dir="/tmp/huggingface")

# -----------------------------
# Tokenizer
# -----------------------------
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

# -----------------------------
# Model Loading
# -----------------------------
device_map = {"": 0}
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    torch_dtype=torch.float16,
    device_map=device_map,
)
model.config.use_cache = False  # Important for memory savings
model.config.pretraining_tp = 1

# -----------------------------
# LoRA Configuration
# -----------------------------
peft_config = LoraConfig(
    lora_alpha=lora_alpha,
    lora_dropout=lora_dropout,
    r=lora_r,
    bias="none",
    task_type="CAUSAL_LM",
)

# -----------------------------
# TrainingArguments
# -----------------------------
training_arguments = TrainingArguments(
    output_dir=output_dir,
    num_train_epochs=1,
    per_device_train_batch_size=4,  # Reduced from 4 to save VRAM
    gradient_accumulation_steps=4,  # Effective batch size = 4
    optim="paged_adamw_32bit",
    save_steps=0,
    logging_steps=25,
    learning_rate=2e-4,
    weight_decay=0.001,
    fp16=True,   # Enabled for memory efficiency
    bf16=False,
    max_grad_norm=0.3,
    max_steps=-1,
    warmup_ratio=0.03,
    group_by_length=True,
    lr_scheduler_type="cosine",
    gradient_checkpointing=True,  # Already True
    push_to_hub=True,  # Disable for now; re-enable after testing
)

# -----------------------------
# Trainer
# -----------------------------
trainer = SFTTrainer(
    model=model,
    args=training_arguments,
    train_dataset=dataset,
    peft_config=peft_config,
)

# -----------------------------
# Train and Save
# -----------------------------
trainer.train()
trainer.save_model(output_dir)