#!/usr/bin/env python3
"""Upload SDG output JSONL to Hugging Face as a dataset.

Usage:
    python -m sdg.upload_to_hf output/my_experiment/output.jsonl
    python -m sdg.upload_to_hf output/my_experiment/output.jsonl --org teetone --private
"""

import argparse
import json
import tempfile
from pathlib import Path

from huggingface_hub import HfApi, create_repo


def load_jsonl(filepath: str) -> list[dict]:
    """Load JSONL file and return list of records."""
    records = []
    with open(filepath, "r") as f:
        for line in f:
            if line.strip():
                records.append(json.loads(line))
    return records


def main():
    parser = argparse.ArgumentParser(description="Upload SDG output JSONL to Hugging Face")
    parser.add_argument("jsonl_file", help="Path to the JSONL file")
    parser.add_argument("--repo-name", default=None, help="Dataset repo name (default: parent directory name)")
    parser.add_argument("--org", default="teetone", help="HF organization/username")
    parser.add_argument("--private", action="store_true", help="Make dataset private")
    args = parser.parse_args()

    jsonl_path = Path(args.jsonl_file)
    if not jsonl_path.exists():
        raise FileNotFoundError(f"File not found: {jsonl_path}")

    # Derive dataset name from parent directory (experiment_name) or filename
    dataset_name = args.repo_name or jsonl_path.parent.name
    repo_id = f"{args.org}/{dataset_name}"

    print(f"Loading {jsonl_path}...")
    records = load_jsonl(str(jsonl_path))
    print(f"Loaded {len(records):,} records")

    # Write clean JSONL with just messages
    print("Preparing data...")
    with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f:
        temp_path = f.name
        for r in records:
            clean_record = {"messages": r["messages"]}
            f.write(json.dumps(clean_record, ensure_ascii=False) + "\n")

    api = HfApi()

    print(f"Creating repo {repo_id}...")
    create_repo(repo_id, repo_type="dataset", private=args.private, exist_ok=True)

    print(f"Uploading to {repo_id}...")
    api.upload_file(
        path_or_fileobj=temp_path,
        path_in_repo="data/train-00000-of-00001.jsonl",
        repo_id=repo_id,
        repo_type="dataset",
    )

    # Also upload stats.json if it exists alongside the JSONL
    stats_path = jsonl_path.parent / "stats.json"
    if stats_path.exists():
        print(f"Uploading {stats_path}...")
        api.upload_file(
            path_or_fileobj=str(stats_path),
            path_in_repo="stats.json",
            repo_id=repo_id,
            repo_type="dataset",
        )

    Path(temp_path).unlink()

    print(f"Done! Dataset available at: https://huggingface.co/datasets/{repo_id}")


if __name__ == "__main__":
    main()