"""
Add Complete AI4Bharat Samanantar (10.1M English-Hindi parallel translation pairs, 1.83 GB)
to ViuAI/viu-mini-raw-pretrain Hub via server-side CommitOperationCopy.
"""
import sys
import time
from huggingface_hub import HfApi, CommitOperationCopy

sys.stdout.reconfigure(encoding='utf-8')

def main():
    api = HfApi()
    source_repo = "ai4bharat/samanantar"
    target_repo = "ViuAI/viu-mini-raw-pretrain"
    subfolder = "hi"
    target_folder = "translation"

    print("=== Adding AI4Bharat Samanantar to Translation Corpus ===")
    print(f"Source: {source_repo} ({subfolder})")
    print(f"Target: {target_repo}/{target_folder}/")

    # List all files in hi/
    repo_files = api.list_repo_files(source_repo, repo_type="dataset")
    sam_files = sorted([f for f in repo_files if f.startswith(f"{subfolder}/") and f.endswith(".parquet")])
    print(f"Total Samanantar Hindi Parquet shards found: {len(sam_files)}")

    target_files = set(api.list_repo_files(target_repo, repo_type="dataset"))
    
    operations = []
    skipped = 0
    for src_path in sam_files:
        filename = src_path.split("/")[-1]
        dst_path = f"{target_folder}/samanantar-{filename}"
        if dst_path in target_files:
            print(f"  [Already Present] {dst_path}")
            skipped += 1
            continue
        operations.append(
            CommitOperationCopy(
                src_path_in_repo=src_path,
                path_in_repo=dst_path,
                src_repo_id=source_repo,
                src_repo_type="dataset"
            )
        )

    print(f"\nSkipped: {skipped} files | To Copy: {len(operations)} files")
    if not operations:
        print("[ok] All Samanantar shards are already present in target repo!")
        return

    print(f"Committing {len(operations)} server-side copies to {target_repo}...")
    start_t = time.time()
    api.create_commit(
        repo_id=target_repo,
        repo_type="dataset",
        operations=operations,
        commit_message="feat(translation): ingest AI4Bharat Samanantar 10.1M parallel English-Hindi sentence pairs (1.83 GB)"
    )
    elapsed = time.time() - start_t
    print(f"[SUCCESS] All {len(operations)} Samanantar shards committed in {elapsed:.2f}s!")

if __name__ == "__main__":
    main()