#!/usr/bin/env python3
"""
prepare_data.py
Downloads datasets from HuggingFace and converts to JOY format.
Runs once during Docker build.
"""

import json
import os
import sys

OUTPUT = "data/master_data.jsonl"
os.makedirs("data", exist_ok=True)

total = 0

def write(f, inp, out):
    global total
    inp = inp.strip().replace('"', "'")
    out = out.strip().replace('"', "'")
    if len(inp) < 3 or len(out) < 3: return
    if len(inp) > 2000 or len(out) > 2000: return
    # remove newlines
    inp = inp.replace('\n', ' ').replace('\r', '')
    out = out.replace('\n', ' ').replace('\r', '')
    f.write(f'{{"input":"{inp}","output":"{out}"}}\n')
    total += 1
    if total % 10000 == 0:
        print(f"[DATA] Written {total} examples...", flush=True)

print("[DATA] Starting dataset preparation...", flush=True)

with open(OUTPUT, "w") as f:

    # ── Dataset 1: Alpaca (52K instruction following) ──
    try:
        print("[DATA] Downloading alpaca dataset...", flush=True)
        from datasets import load_dataset
        ds = load_dataset("tatsu-lab/alpaca", split="train")
        for row in ds:
            inp = row.get("instruction", "")
            ctx = row.get("input", "")
            out = row.get("output", "")
            if ctx: inp = f"{inp} {ctx}"
            write(f, inp, out)
        print(f"[DATA] Alpaca done. Total: {total}", flush=True)
    except Exception as e:
        print(f"[DATA] Alpaca failed: {e}", flush=True)

    # ── Dataset 2: Dolly 15K ──
    try:
        print("[DATA] Downloading dolly dataset...", flush=True)
        ds = load_dataset("databricks/databricks-dolly-15k", split="train")
        for row in ds:
            inp = row.get("instruction", "")
            ctx = row.get("context", "")
            out = row.get("response", "")
            if ctx: inp = f"{inp} Context: {ctx}"
            write(f, inp, out)
        print(f"[DATA] Dolly done. Total: {total}", flush=True)
    except Exception as e:
        print(f"[DATA] Dolly failed: {e}", flush=True)

    # ── Dataset 3: OpenAssistant ──
    try:
        print("[DATA] Downloading OpenAssistant...", flush=True)
        ds = load_dataset("OpenAssistant/oasst1", split="train")
        # pair prompts with responses
        messages = {}
        for row in ds:
            mid = row.get("message_id")
            pid = row.get("parent_id")
            role = row.get("role")
            text = row.get("text", "")
            messages[mid] = {"role": role, "text": text, "parent": pid}

        for mid, msg in messages.items():
            if msg["role"] == "assistant" and msg["parent"]:
                parent = messages.get(msg["parent"])
                if parent and parent["role"] == "prompter":
                    write(f, parent["text"], msg["text"])
        print(f"[DATA] OA done. Total: {total}", flush=True)
    except Exception as e:
        print(f"[DATA] OA failed: {e}", flush=True)

    # ── Dataset 4: Natural Questions ──
    try:
        print("[DATA] Downloading NQ...", flush=True)
        ds = load_dataset("nq-open", split="train")
        for row in ds:
            q = row.get("question", "")
            a = row.get("answer", [])
            if a and isinstance(a, list): a = a[0]
            if isinstance(a, str): write(f, q, a)
        print(f"[DATA] NQ done. Total: {total}", flush=True)
    except Exception as e:
        print(f"[DATA] NQ failed: {e}", flush=True)

    # ── Dataset 5: SciQ ──
    try:
        print("[DATA] Downloading SciQ...", flush=True)
        ds = load_dataset("allenai/sciq", split="train")
        for row in ds:
            q = row.get("question", "")
            a = row.get("correct_answer", "")
            sup = row.get("support", "")
            if sup: out = f"{a}. {sup}"
            else:   out = a
            write(f, q, out)
        print(f"[DATA] SciQ done. Total: {total}", flush=True)
    except Exception as e:
        print(f"[DATA] SciQ failed: {e}", flush=True)

    # ── Dataset 6: TriviaQA ──
    try:
        print("[DATA] Downloading TriviaQA...", flush=True)
        ds = load_dataset("trivia_qa", "rc.nocontext", split="train")
        for row in ds:
            q = row.get("question", "")
            a = row.get("answer", {})
            if isinstance(a, dict):
                aliases = a.get("aliases", [])
                if aliases: a = aliases[0]
                else: a = a.get("value", "")
            if isinstance(a, str): write(f, q, a)
        print(f"[DATA] TriviaQA done. Total: {total}", flush=True)
    except Exception as e:
        print(f"[DATA] TriviaQA failed: {e}", flush=True)

    # ── Dataset 7: CommonsenseQA ──
    try:
        print("[DATA] Downloading CommonsenseQA...", flush=True)
        ds = load_dataset("tau/commonsense_qa", split="train")
        for row in ds:
            q = row.get("question", "")
            ans_key = row.get("answerKey", "")
            choices = row.get("choices", {})
            labels = choices.get("label", [])
            texts  = choices.get("text", [])
            for label, text in zip(labels, texts):
                if label == ans_key:
                    write(f, q, text)
                    break
        print(f"[DATA] CommonsenseQA done. Total: {total}", flush=True)
    except Exception as e:
        print(f"[DATA] CommonsenseQA failed: {e}", flush=True)

print(f"\n[DATA] ✅ Preparation complete!", flush=True)
print(f"[DATA] Total examples: {total}", flush=True)
print(f"[DATA] Output: {OUTPUT}", flush=True)