import os
import json
import sys
from datasets import load_dataset

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import config

def extract_prompt_chosen_rejected(example):
    prompt = None
    chosen = None
    rejected = None

    if "chosen" in example and "rejected" in example:
        c_list = example["chosen"]
        r_list = example["rejected"]

        if isinstance(c_list, list) and len(c_list) > 0:
            last_c = c_list[-1]
            chosen = last_c.get("content", "") if isinstance(last_c, dict) else str(last_c)
            if len(c_list) > 1:
                prompt_item = c_list[0]
                prompt = prompt_item.get("content", "") if isinstance(prompt_item, dict) else str(prompt_item)

        if isinstance(r_list, list) and len(r_list) > 0:
            last_r = r_list[-1]
            rejected = last_r.get("content", "") if isinstance(last_r, dict) else str(last_r)
            if not prompt and len(r_list) > 1:
                prompt_item = r_list[0]
                prompt = prompt_item.get("content", "") if isinstance(prompt_item, dict) else str(prompt_item)

    if not prompt and "prompt" in example:
        prompt = example["prompt"]

    if prompt and chosen and rejected:
        return {
            "prompt": prompt.strip(),
            "chosen": chosen.strip(),
            "rejected": rejected.strip()
        }
    return None

def main():
    print(f"Downloading RLHF preference dataset: {config.DATASET_ID}...")
    dataset = load_dataset(config.DATASET_ID, split="train")

    print(f"Loaded {len(dataset)} raw samples. Filtering and reformatting DPO preference pairs...")
    formatted_items = []
    
    for example in dataset:
        item = extract_prompt_chosen_rejected(example)
        if item and len(item["prompt"]) > 10 and len(item["chosen"]) > 20 and len(item["rejected"]) > 20:
            formatted_items.append(item)
            if len(formatted_items) >= 10000:
                break

    print(f"Successfully processed {len(formatted_items)} valid preference pairs.")

    val_size = min(1000, int(len(formatted_items) * 0.1))
    train_size = len(formatted_items) - val_size

    train_items = formatted_items[:train_size]
    val_items = formatted_items[train_size:]

    os.makedirs(config.DATA_DIR, exist_ok=True)

    print(f"Saving {len(train_items)} training pairs to {config.TRAIN_FILE}...")
    with open(config.TRAIN_FILE, "w", encoding="utf-8") as f:
        for item in train_items:
            f.write(json.dumps(item) + "\n")

    print(f"Saving {len(val_items)} validation pairs to {config.VAL_FILE}...")
    with open(config.VAL_FILE, "w", encoding="utf-8") as f:
        for item in val_items:
            f.write(json.dumps(item) + "\n")

    print("RLHF Preference dataset preparation complete!")

if __name__ == "__main__":
    main()