import torch
import wandb
from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import GRPOConfig, GRPOTrainer
from datasets import load_dataset, Dataset
from huggingface_hub import login
from common.constants import SYSTEM_PROMPT
from common.rewards import SmoothXMLRewardEvaluator

##########
# CONFIG #
##########

MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
DATASET = "krogoldAI/rag-ambiguous-queries"
OUTPUT_MODEL = "model/repo"
HF_TOKEN = "<MY_HF_TOKEN>"

# Training hyperparameters
NUM_EPOCHS = 1
PER_DEVICE_BATCH = 8
GRADIENT_ACCUMULATION = 1 # 2
LEARNING_RATE = 5e-6 # Low for GRPO stability
NUM_GENERATIONS = 2
MAX_NEW_TOKENS = 512
WARMUP_STEPS = 50

#########
# LOGIN #
#########

print("Logging into Hugging Face...")
login(HF_TOKEN)

##################
# LOAD TOKENIZER #
##################

print(f"Loading tokenizer from {MODEL_NAME}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.padding_side = "left"

# Ensure pad token is set
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

################
# LOAD DATASET #
################

print(f"Loading dataset {DATASET}...")
raw_dataset = load_dataset(DATASET, "default", split="train")

# Extract the list of questions
questions_list = raw_dataset["questions"]
print(f"Dataset size: {len(questions_list)} examples")

# Create a proper dataset from the list
dataset = Dataset.from_dict({"question": questions_list})
dataset = dataset.shuffle(seed=42)

# Take on 3000 examples
dataset = dataset.select(range(3000))

###########################
# PREPROCESSING FUNCTION  #
###########################

def preprocess(example):
    """
    Apply chat template to create proper prompt format.
    """
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": example["question"]},
    ]
    
    # Apply chat template without tokenization to get the prompt string
    prompt = tokenizer.apply_chat_template(
        messages, 
        tokenize=False, 
        add_generation_prompt=True
    )
    
    return {
        "prompt": prompt,
        "query": example["question"]  # Keep for reference
    }

print("Preprocessing dataset with chat template...")
dataset = dataset.map(preprocess, remove_columns=dataset.column_names)

##############
# LOAD MODEL #
##############

print(f"\nLoading model {MODEL_NAME}...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="flash_attention_2",
    use_cache=False,
)

# If Liger Kernel available, apply now
try:
    from liger_kernel.transformers import apply_liger_kernel_to_qwen2
    try:
        apply_liger_kernel_to_qwen2(model)
    except TypeError:
        apply_liger_kernel_to_qwen2()
    print("Liger Kernel applied successfully for Qwen2 optimization")
except Exception as e:
    print(f"Liger Kernel not available or failed to apply: {e}")

print(f"Model loaded. Parameters: {model.num_parameters() / 1e9:.3f}B")

model.generation_config.max_new_tokens = MAX_NEW_TOKENS
model.generation_config.temperature = 0.7
model.generation_config.do_sample = True

###################
# REWARD FUNCTION #
###################

# Initialize reward evaluator
evaluator = SmoothXMLRewardEvaluator()

def reward_function(completions, prompts, **kwargs):
    rewards = []
    
    for i, (completion, prompt) in enumerate(zip(completions, prompts)):
        # Extract generated portion
        if completion.startswith(prompt):
            response = completion[len(prompt):].strip()
        else:
            response = completion        
        try:
            analysis = evaluator.get_detailed_analysis(response)
            score = analysis["final_score"]            
            rewards.append(score)
            
        except Exception as e:
            print(f"ERROR in reward evaluation: {e}")
            print(f"Response that failed: {response[:200]}")
            rewards.append(0.0)
    
    return rewards


##############
# MONITORING #
##############

# Monitoring the LLM from WandB
wandb.login(key = "<MY_WANDB_KEY>")
run = wandb.init(project='QueryRefiner-GRPO-Phase1', job_type="training", anonymous="allow")

######################
# TRAINING ARGUMENTS #
######################

print("Configuring training arguments...")
training_args = GRPOConfig(
    output_dir="./qwen_grpo_checkpoints",
    run_name="qwen-grpo-query-analysis",
    num_train_epochs=NUM_EPOCHS,
    per_device_train_batch_size=PER_DEVICE_BATCH,
    gradient_accumulation_steps=GRADIENT_ACCUMULATION,
    learning_rate=LEARNING_RATE,
    num_generations=NUM_GENERATIONS,
    optim="adamw_torch_fused",
    warmup_steps=WARMUP_STEPS,
    lr_scheduler_type="cosine",
    weight_decay=0.01,
    max_grad_norm=1.0,
    bf16=True,
    gradient_checkpointing=True,
    use_liger_kernel=True,
    logging_steps=10,
    save_steps=200,
    save_total_limit=3,
    push_to_hub=True,
    hub_model_id=OUTPUT_MODEL,
    hub_private_repo=True,
    hub_strategy="checkpoint",
    dataloader_num_workers=4,
    dataloader_pin_memory=True,
    remove_unused_columns=False,
    report_to="wandb",
    seed=42,
)

###########
# TRAINER #
###########

print("Initializing GRPOTrainer...")
trainer = GRPOTrainer(
    model=model,
    reward_funcs=reward_function,
    args=training_args,
    train_dataset=dataset, # Uses by default the "prompt" column, all other columns are ignored
)

#########
# TRAIN #
#########

print("=" * 50)
print("Starting GRPO training...")
print(f"Model: {MODEL_NAME}")
print(f"Dataset size: {len(dataset)}")
print(f"Epochs: {NUM_EPOCHS}")
print(f"Effective batch size: {PER_DEVICE_BATCH * GRADIENT_ACCUMULATION}")
print(f"Generations per prompt: {NUM_GENERATIONS}")
print("=" * 50)

trainer.train()

####################
# SAVE FINAL MODEL #
####################

print("\nSaving final model...")
final_model = trainer.model
final_model.config.use_cache = True  # Re-enable cache for inference
final_model.save_pretrained("./final_grpo_model")
tokenizer.save_pretrained("./final_grpo_model")

###############
# PUSH TO HUB #
###############

print(f"\nPushing final model to {OUTPUT_MODEL}...")
try:
    final_model.push_to_hub(OUTPUT_MODEL, private=True)
    tokenizer.push_to_hub(OUTPUT_MODEL, private=True)
    print("Successfully pushed model and tokenizer to Hub!")
except Exception as e:
    print(f"Warning: Failed to push to hub: {e}")
    print("Model saved locally at ./final_grpo_model")

print("\n" + "=" * 50)
print("GRPO Training Complete!")
print("=" * 50)