import json
import os
from pathlib import Path

from datasets import load_dataset
from dotenv import load_dotenv

load_dotenv()

OUTPUT_DIR = Path("data/raw_sources/hackaprompt")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

OUTPUT_FILE = OUTPUT_DIR / "hackaprompt_full.jsonl"

def main():
    # Replace this with the exact dataset ID your team uses on Hugging Face
    dataset_id = "hackaprompt/hackaprompt-dataset"

    token = os.getenv("HUGGINGFACE_TOKEN") or os.getenv("HF_TOKEN")
    load_kw = {"split": "train"}
    if token:
        load_kw["token"] = token
    ds = load_dataset(dataset_id, **load_kw)

    with OUTPUT_FILE.open("w", encoding="utf-8") as f:
        for row in ds:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")

    print(f"Saved full HackAPrompt dataset to: {OUTPUT_FILE}")
    print(f"Total rows written: {len(ds)}")

if __name__ == "__main__":
    main()