#!/usr/bin/env python3
"""Publish the AI Puppet Theater Actor SFT dataset to Hugging Face Hub.

Dry-run mode performs no network calls and is safe to run without HF auth.
"""

from __future__ import annotations

import argparse
import os
from pathlib import Path


DEFAULT_REPO_ID = "build-small-hackathon/AI-Puppet-Theater-Actor-SFT"
DEFAULT_DATASET_DIR = Path("finetune/data")
DEFAULT_SAMPLE_DIR = Path("finetune/data_samples")
DEFAULT_CARD = Path("finetune/dataset_cards/actor_sft_README.md")
DATASET_FILES = [
    "actor_sft_v0.jsonl",
    "actor_sft_v0_train.jsonl",
    "actor_sft_v0_val.jsonl",
    "actor_sft_v1.jsonl",
    "actor_sft_v1_train.jsonl",
    "actor_sft_v1_val.jsonl",
]
SAMPLE_FILES = [
    "actor_sft_v0_sample.jsonl",
    "actor_sft_v1_sample.jsonl",
    "actor_eval_prompts.jsonl",
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo_id", default=DEFAULT_REPO_ID)
    parser.add_argument("--dataset_dir", type=Path, default=DEFAULT_DATASET_DIR)
    parser.add_argument("--sample_dir", type=Path, default=DEFAULT_SAMPLE_DIR)
    parser.add_argument("--card", type=Path, default=DEFAULT_CARD)
    parser.add_argument("--private", action="store_true")
    parser.add_argument("--dry_run", action=argparse.BooleanOptionalAction, default=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    files = collect_files(args.dataset_dir, args.sample_dir, args.card)
    if args.dry_run:
        print_plan(args.repo_id, args.private, files)
        return
    publish_dataset(args.repo_id, args.private, files)


def collect_files(dataset_dir: Path, sample_dir: Path, card: Path) -> list[tuple[Path, str]]:
    files: list[tuple[Path, str]] = []
    if card.exists():
        files.append((card, "README.md"))
    else:
        print(f"warning: dataset card missing: {card}")
    for name in DATASET_FILES:
        path = dataset_dir / name
        if path.exists():
            files.append((path, f"data/{name}"))
        else:
            print(f"warning: dataset file missing: {path}")
    for name in SAMPLE_FILES:
        path = sample_dir / name
        if path.exists():
            files.append((path, f"data_samples/{name}"))
        else:
            print(f"warning: sample file missing: {path}")
    return files


def print_plan(repo_id: str, private: bool, files: list[tuple[Path, str]]) -> None:
    print("DRY RUN: no Hugging Face network calls or uploads will be performed.")
    print(f"repo_type: dataset")
    print(f"repo_id: {repo_id}")
    print(f"private: {private}")
    print("files:")
    for local_path, repo_path in files:
        print(f"  {local_path} -> {repo_path}")


def publish_dataset(repo_id: str, private: bool, files: list[tuple[Path, str]]) -> None:
    token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
    if not token:
        raise SystemExit("HF_TOKEN or HUGGINGFACEHUB_API_TOKEN is required for a real upload. Re-run with --dry_run to inspect the plan.")
    if not files:
        raise SystemExit("No files found to upload.")
    from huggingface_hub import HfApi

    api = HfApi(token=token)
    api.create_repo(repo_id=repo_id, repo_type="dataset", private=private, exist_ok=True)
    for local_path, repo_path in files:
        api.upload_file(
            path_or_fileobj=str(local_path),
            path_in_repo=repo_path,
            repo_id=repo_id,
            repo_type="dataset",
        )
        print(f"uploaded {local_path} -> {repo_path}")
    print(f"dataset publish complete: https://huggingface.co/datasets/{repo_id}")


if __name__ == "__main__":
    main()