"""
Push LUNA training code + config to HuggingFace.

Code repo:  https://huggingface.co/spaces/ASTERIZER/LUNA
Model repo: https://huggingface.co/ASTERIZER/LUNA-100M (tokenizer only)

Usage:
    HF_TOKEN=hf_xxx python push_code_to_hf.py
    HF_TOKEN=hf_xxx python push_code_to_hf.py --repo ASTERIZER/LUNA --type space
"""

import argparse
import os
from pathlib import Path
from huggingface_hub import HfApi, create_repo

DEFAULT_CODE_REPO = "ASTERIZER/LUNA"
DEFAULT_REPO_TYPE = "space"
TOKEN = os.environ.get("HF_TOKEN")

FILES_TO_PUSH = [
    # Core training scripts
    "train.py",
    "train_300m.py",
    "sft_train.py",
    "lora_sft_train.py",
    "chat.py",
    "chat_full_sft.py",
    "generate.py",

    # Configs
    "train_config.yaml",
    "train_config_300m.yaml",
    "train_continue_english_1b.yaml",
    "sft_config.yaml",
    "rag_mcp_lora_config.yaml",
    "rag_mcp_full_sft_config.yaml",

    # Data pipeline scripts
    "Base/scripts/build_english_1b.py",
    "Base/scripts/clean_english_1b.py",
    "Base/scripts/prepare_litdata.py",
    "Base/scripts/build_english_corpus.py",
    "Base/scripts/build_english_curriculum_1b.py",
    "Base/scripts/build_instruct_dataset.py",
    "Base/scripts/filter_datasets.py",
    "Base/scripts/deep_clean_sft.py",

    # HF upload / push scripts
    "push_code_to_hf.py",
    "push_dataset_to_hf.py",
    "push_model_to_hf.py",
    "upload_lora_to_hf.py",
    "upload_full_sft_to_hf.py",

    # Validation / benchmarking
    "validate_sft.py",
    "check_sft_alignment.py",
    "validate_and_quantize.py",
    "benchmark_runpod.py",

    # Smoke test
    "smoke_test_300m.py",

    # Instance run scripts
    "run_cloud_300m.sh",
    "run_english_1b_instance.sh",
    "setup_and_train.sh",
    "setup_and_train_300m.sh",
    "setup_and_sft.sh",
    "gpu_train.sh",
    "gpu_full_sft.sh",

    # Requirements & docs
    "requirements.txt",
    "README.md",
    "fetch_data.py",

    # Dataset-related
    "Base/Datasets/rag_mcp_sft/build_rag_mcp_sft_dataset.py",
    "Base/Datasets/rag_mcp_sft/push_to_hf.py",
    "Base/Datasets/rag_mcp_sft/BUILD_REPORT.md",
    "Base/Datasets/rag_mcp_sft/FINETUNE_COMMANDS.md",
    "Base/Datasets/rag_mcp_sft/README.md",
    "Base/Datasets/rag_mcp_sft/source_manifest.json",
    "Base/Datasets/rag_mcp_sft/sample_preview.json",

    # Tokenizer config (small files only)
    "Base/checkpoints/EleutherAI/pythia-160m/config.json",
    "Base/checkpoints/EleutherAI/pythia-160m/tokenizer_config.json",
    "Base/checkpoints/EleutherAI/pythia-160m/tokenizer.json",
]

EVALUATION_GLOBS = [
    "Evaluation/*.py",
    "Evaluation/*.sh",
    "Evaluation/*.yaml",
    "Evaluation/*.yml",
    "Evaluation/*.md",
    "Evaluation/*.txt",
    "Evaluation/*.json",
]


def build_file_list():
    files = []
    seen = set()

    for fpath in FILES_TO_PUSH:
        if fpath not in seen:
            files.append(fpath)
            seen.add(fpath)

    for pattern in EVALUATION_GLOBS:
        for path in sorted(Path(".").glob(pattern)):
            if not path.is_file():
                continue
            rel = path.as_posix()
            if rel not in seen:
                files.append(rel)
                seen.add(rel)

    return files


def main():
    parser = argparse.ArgumentParser(description="Push LUNA code to HuggingFace")
    parser.add_argument("--repo", default=DEFAULT_CODE_REPO, help="HF repo ID")
    parser.add_argument("--type", default=DEFAULT_REPO_TYPE,
                        choices=["space", "model"], help="Repo type")
    args = parser.parse_args()

    token = TOKEN
    if not token:
        raise RuntimeError("Set HF_TOKEN environment variable")

    api = HfApi(token=token)

    create_repo(
        repo_id=args.repo,
        token=token,
        repo_type=args.type,
        exist_ok=True,
        private=False,
        space_sdk="static" if args.type == "space" else None,
    )
    print(f"Repo ready: https://huggingface.co/{'spaces/' if args.type == 'space' else ''}{args.repo}")

    files_to_push = build_file_list()
    print(f"Preparing to push {len(files_to_push)} files...")

    pushed = 0
    for fpath in files_to_push:
        if not os.path.exists(fpath):
            print(f"  SKIP (not found): {fpath}")
            continue
        api.upload_file(
            path_or_fileobj=fpath,
            path_in_repo=fpath,
            repo_id=args.repo,
            repo_type=args.type,
            token=token,
        )
        print(f"  OK: {fpath}")
        pushed += 1

    print(f"\nPushed {pushed}/{len(files_to_push)} files to https://huggingface.co/{'spaces/' if args.type == 'space' else ''}{args.repo}")


if __name__ == "__main__":
    main()