import shutil
from pathlib import Path
import sys

# Experiment directory to be used for deployment is passed in as input argument
exp_dir = Path(sys.argv[1])
deploy_dir = Path("hf_deployment")

# Safety-first: enables first seeing the changes before actually transfering files over (AWS s3 operations-like)
dry_run = "--dry-run" in sys.argv
checkpoint = deploy_dir / "checkpoint"
MODEL_SUBDIR = "model"
prefix = "[DRY RUN]" if dry_run else ""

print(f"{prefix} Create: {checkpoint}")
if not dry_run:
    checkpoint.mkdir(parents=True, exist_ok=True)
    (checkpoint / MODEL_SUBDIR).mkdir(exist_ok=True)

# Individual files
files = [
    (exp_dir / "config.json", checkpoint / "config.json"),
    (
        exp_dir / MODEL_SUBDIR / "pytorch_model.bin",
        checkpoint / MODEL_SUBDIR / "pytorch_model.bin",
    ),
]
for src, dst in files:
    print(f"{prefix} Copy: {src} -> {dst}")
    if not dry_run:
        shutil.copy2(src, dst)

# Directories (recursively)
trees = [
    (exp_dir / MODEL_SUBDIR / "tokenizer", checkpoint / MODEL_SUBDIR / "tokenizer"),
    (Path("src"), deploy_dir / "src"),
]
for src, dst in trees:
    print(f"{prefix} Copy tree: {src} -> {dst}")
    if not dry_run:
        shutil.copytree(src, dst, dirs_exist_ok=True)

if not dry_run:
    print(f"\nDeployment files are ready under {deploy_dir}.")