import os
import torch
from datasets import load_dataset
from transformers import (
    AutoTokenizer, 
    AutoModelForCausalLM, 
    TrainingArguments, 
    Trainer, 
    DataCollatorForLanguageModeling, 
    TrainerCallback
)
from peft import LoraConfig, get_peft_model
import spaces
import time

# Đường dẫn lưu checkpoint
CHECKPOINT_DIR = "./checkpoints"
if not os.path.exists(CHECKPOINT_DIR):
    os.makedirs(CHECKPOINT_DIR)

# Tải Dataset (CPU)
dataset = load_dataset('vntc/wiki-mini-corpus')

# Chia Dataset thành train và validation (CPU)
split_dataset = dataset['train'].train_test_split(test_size=0.1, seed=42)
train_dataset = split_dataset['train']
validation_dataset = split_dataset['test']

# Tiền Xử Lý Văn Bản (CPU)
def preprocess_function(examples):
    passages = [passage.lower().strip() for passage in examples['passage']]
    return {'passage': passages}

processed_train = train_dataset.map(preprocess_function, batched=True, remove_columns=['id', 'metadata'])
processed_validation = validation_dataset.map(preprocess_function, batched=True, remove_columns=['id', 'metadata'])

# Tokenization (CPU)
model_name = "meta-llama/Llama-3.2-3B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Đảm bảo tokenizer có pad_token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

def tokenize_function(examples):
    return tokenizer(
        examples['passage'],
        padding='max_length',
        truncation=True,
        max_length=512,
    )

tokenized_train = processed_train.map(tokenize_function, batched=True)
tokenized_validation = processed_validation.map(tokenize_function, batched=True)

# Thêm trường 'labels' (CPU)
def add_labels(examples):
    examples['labels'] = examples['input_ids'].copy()
    return examples

tokenized_train = tokenized_train.map(add_labels, batched=True)
tokenized_validation = tokenized_validation.map(add_labels, batched=True)

# Loại bỏ các cột không cần thiết (CPU)
tokenized_train = tokenized_train.remove_columns(['passage'])
tokenized_validation = tokenized_validation.remove_columns(['passage'])

# Định dạng dữ liệu cho PyTorch (CPU)
tokenized_train.set_format('torch')
tokenized_validation.set_format('torch')

# Tạo DatasetDict (CPU)
final_dataset = {
    'train': tokenized_train,
    'validation': tokenized_validation
}

# Định Nghĩa TrainerCallback để Lưu Checkpoint Nhanh Hơn
class SaveCheckpointCallback(TrainerCallback):
    def on_step_end(self, args, state, control, **kwargs):
        if state.global_step % args.save_steps == 0 and state.global_step != 0:
            checkpoint_path = os.path.join(args.output_dir, f"checkpoint-{state.global_step}")
            print(f"Lưu checkpoint tại: {checkpoint_path}")
            trainer = kwargs['trainer']  # Truy cập trainer từ kwargs
            trainer.save_model(checkpoint_path)
        return control  # Trả về đối tượng control hiện tại

# Định Nghĩa Hàm Huấn Luyện với Decorator @spaces.GPU
@spaces.GPU(duration=15, queue=False)
def run_training():
    """
    Hàm huấn luyện mô hình sử dụng GPU với thời gian hạn chế.
    """
    # Tải và Cấu Hình Mô Hình với LoRA (GPU)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        device_map="auto",
        torch_dtype=torch.float16,
        load_in_8bit=False
    )
    
    lora_config = LoraConfig(
        r=8,
        lora_alpha=32,
        target_modules=["q_proj", "k_proj", "v_proj", "out_proj"],
        lora_dropout=0.1,
        bias="none",
    )
    
    model = get_peft_model(model, lora_config)
    print(model)
    
    # Cấu Hình TrainingArguments (GPU)
    training_args = TrainingArguments(
        output_dir=CHECKPOINT_DIR,
        per_device_train_batch_size=4,
        per_device_eval_batch_size=4,
        gradient_accumulation_steps=8,
        num_train_epochs=3,
        max_steps=50,  # Đặt max_steps tại đây
        learning_rate=3e-4,
        weight_decay=0.01,
        logging_steps=5,  # Giảm số bước logging để theo dõi thường xuyên hơn
        eval_strategy="steps",  # Đánh giá sau mỗi vài bước
        eval_steps=5,  # Đánh giá sau mỗi 50 bước
        save_strategy="steps",  # Lưu checkpoint sau mỗi vài bước
        save_steps=5,  # Lưu checkpoint sau mỗi 50 bước
        save_total_limit=5,  # Giới hạn số lượng checkpoint lưu trữ
        fp16=True,
        report_to="none",
        load_best_model_at_end=True,
    )
    
    # Data Collator (GPU)
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer, 
        mlm=False,  # Vì bạn đang thực hiện Causal LM
        pad_to_multiple_of=8
    )
    
    # Tạo Trainer (GPU)
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=final_dataset['train'],
        eval_dataset=final_dataset['validation'],
        tokenizer=tokenizer,
        data_collator=data_collator,
        callbacks=[SaveCheckpointCallback()],  # Thêm callback
    )
    
    # Kiểm tra nếu có checkpoint
    checkpoints = [os.path.join(CHECKPOINT_DIR, d) for d in os.listdir(CHECKPOINT_DIR) if d.startswith('checkpoint')]
    if checkpoints:
        latest_checkpoint = max(checkpoints, key=os.path.getctime)
        print(f"Đang tiếp tục huấn luyện từ checkpoint: {latest_checkpoint}")
        trainer.train(resume_from_checkpoint=latest_checkpoint)
    else:
        trainer.train()
    
    # Lưu checkpoint sau khi huấn luyện
    trainer.save_model(CHECKPOINT_DIR)
    return "Huấn luyện hoàn tất hoặc đã tiếp tục từ checkpoint."

# Hàm Tự Động Hóa Việc Gọi Lặp Lại Hàm Huấn Luyện
def continuous_training(total_steps=300, steps_per_call=5):
    """
    Hàm tự động gọi lại `run_training` để hoàn thành quá trình huấn luyện.
    
    Args:
        total_steps (int): Tổng số bước huấn luyện mong muốn.
        steps_per_call (int): Số bước huấn luyện mỗi lần gọi hàm.
    """
    steps_done = 0
    while steps_done < total_steps:
        print(f"Bắt đầu huấn luyện cho {steps_per_call} bước.")
        result = run_training()
        print(result)
        steps_done += steps_per_call
        print(f"Đã huấn luyện {steps_done} / {total_steps} bước.")
        
        # Kiểm tra nếu đã đạt số bước mong muốn
        if steps_done >= total_steps:
            print("Đã hoàn thành toàn bộ quá trình huấn luyện.")
            break
        
        # Chờ một khoảng thời gian trước khi gọi lại (tùy thuộc vào yêu cầu của hệ thống)
        time.sleep(2)  # Thời gian chờ có thể điều chỉnh

# # Gọi hàm huấn luyện liên tục
# if __name__ == "__main__":
#     continuous_training(total_steps=300, steps_per_call=10)