"""HF Job worker: duplicates a Hugging Face repo or bucket.

Strategy by --type:
  - bucket                   : /data/src and /data/dst are mounted as Volumes
                               by the Jobs platform; we copy file by file (lazy).
  - model | dataset | space  : git mirror clone + push, which preserves the
                               full git history (commits, branches, tags) and
                               all LFS objects. The destination is pre-created
                               by the Space with `region=` / `space_sdk=`.

Progress is printed to stdout so the calling Space can tail it via `fetch_job_logs`.
"""

from __future__ import annotations

import argparse
import os
import platform
import shutil
import subprocess
import sys
import tempfile
import time
import urllib.request
import zipfile
from pathlib import Path

CHUNK_SIZE = 64 * 1024 * 1024  # 64 MiB

BUCKET_SRC_MOUNT = "/data/src"
BUCKET_DST_MOUNT = "/data/dst"

HF_ENDPOINT = os.environ.get("HF_ENDPOINT", "https://huggingface.co").rstrip("/")

# URL prefix per repo type. Models live at the root; datasets and spaces have
# their own sub-paths in HF's git layout.
REPO_TYPE_URL_PREFIX = {
    "model": "",
    "dataset": "datasets/",
    "space": "spaces/",
}

# Pin a specific git-xet release so the build is reproducible. Override with
# the GIT_XET_VERSION env var (e.g. "v0.2.2") if needed.
GIT_XET_VERSION = os.environ.get("GIT_XET_VERSION", "v0.2.1")


def log(msg: str) -> None:
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)


def _run(cmd: list[str], cwd: str | None = None) -> None:
    """Run a command, streaming output and raising on non-zero exit."""
    subprocess.run(cmd, cwd=cwd, check=True)


def _run_capture(cmd: list[str], cwd: str | None = None) -> str:
    """Run a command and return stdout, raising on non-zero exit."""
    return subprocess.run(cmd, cwd=cwd, check=True, capture_output=True, text=True).stdout


def copy_bucket() -> int:
    src = Path(BUCKET_SRC_MOUNT)
    dst = Path(BUCKET_DST_MOUNT)

    if not src.is_dir():
        log(f"source mount not found at {src}")
        return 1
    if not dst.is_dir():
        log(f"destination mount not found at {dst}")
        return 1

    log(f"scanning {src} ...")
    files = [f for f in src.rglob("*") if f.is_file()]
    total = len(files)
    bytes_total = sum(f.stat().st_size for f in files)
    log(f"found {total} file(s), {bytes_total / 1e9:.2f} GB")

    if total == 0:
        log("nothing to copy.")
        return 0

    bytes_copied = 0
    failures = 0
    started = time.monotonic()

    for i, src_file in enumerate(files, 1):
        rel = src_file.relative_to(src)
        dst_file = dst / rel
        size = src_file.stat().st_size
        log(f"[{i}/{total}] {rel} ({size / 1e6:.1f} MB)")

        try:
            dst_file.parent.mkdir(parents=True, exist_ok=True)
            with open(src_file, "rb") as sf, open(dst_file, "wb") as df:
                while True:
                    data = sf.read(CHUNK_SIZE)
                    if not data:
                        break
                    df.write(data)
        except Exception as e:
            log(f"  error copying {rel}: {e}")
            failures += 1
            continue

        bytes_copied += size
        pct = bytes_copied / bytes_total * 100 if bytes_total else 100.0
        elapsed = time.monotonic() - started
        rate = bytes_copied / elapsed / 1e6 if elapsed > 0 else 0
        log(f"  done ({pct:.1f}% overall, {rate:.1f} MB/s avg)")

    elapsed = time.monotonic() - started
    log(
        f"transfer complete: {total - failures}/{total} file(s), "
        f"{bytes_copied / 1e9:.2f} GB in {elapsed:.0f}s "
        f"({(bytes_copied / elapsed / 1e6) if elapsed > 0 else 0:.1f} MB/s)"
    )
    if failures:
        log(f"{failures} file(s) failed — see errors above")
        return 2
    return 0


def _ensure_git_lfs() -> None:
    """Install git-lfs at runtime if it isn't already available."""
    if shutil.which("git-lfs") is not None:
        return
    log("installing git-lfs ...")
    _run(["apt-get", "update", "-qq"])
    _run(["apt-get", "install", "-yq", "--no-install-recommends", "git-lfs", "ca-certificates"])
    _run(["git", "lfs", "install", "--system"])


def _ensure_git_xet_binary() -> bool:
    """Drop the git-xet binary into /usr/local/bin without touching git config.

    The upstream installer runs `git xet install`, which writes
    `lfs.customtransfer.xet` into the user's global git config. With that in
    place, `git clone --mirror` of a xet-backed source dies with
    `fatal: expected 'packfile'` (the customtransfer process appears to
    interfere with the smart-HTTP pack negotiation).

    We split the install in two: download the binary here (no config side
    effects), and enable xet locally in the cloned repo only — see
    `_enable_xet_for_push` — so the clone runs in plain-git mode and only
    the push phase routes xet-backed blobs through git-xet.

    Returns True if the binary is available on PATH, False otherwise.
    """
    if shutil.which("git-xet") is not None:
        return True

    machine = platform.machine().lower()
    if machine in ("x86_64", "amd64"):
        arch = "x86_64"
    elif machine in ("aarch64", "arm64"):
        arch = "aarch64"
    else:
        log(f"unsupported architecture for git-xet ({machine}); skipping xet support")
        return False

    url = (
        f"https://github.com/huggingface/xet-core/releases/download/"
        f"git-xet-{GIT_XET_VERSION}/git-xet-linux-{arch}.zip"
    )
    log(f"downloading git-xet {GIT_XET_VERSION} (binary only, no global install) ...")

    with tempfile.TemporaryDirectory() as tmp:
        zip_path = os.path.join(tmp, "git-xet.zip")
        urllib.request.urlretrieve(url, zip_path)
        with zipfile.ZipFile(zip_path) as z:
            z.extractall(tmp)
        bin_path = os.path.join(tmp, "git-xet")
        if not os.path.isfile(bin_path):
            log(f"git-xet binary not found inside {url}; skipping xet support")
            return False
        os.chmod(bin_path, 0o755)
        shutil.move(bin_path, "/usr/local/bin/git-xet")

    return True


def _enable_xet_for_push(clone_dir: str) -> None:
    """Register git-xet as an LFS custom-transfer agent in the LOCAL repo
    config only (no global side effects).

    During `git lfs push`, the LFS client asks the destination which transfer
    protocols it supports. If the destination is xet-backed, it announces
    `xet` — our client then picks the matching customtransfer and spawns
    `git-xet transfer`, which uploads chunks via the xet protocol (parallel,
    content-addressed) instead of the LFS basic blob protocol. For non-xet
    destinations the customtransfer stays dormant; LFS falls back to `basic`.

    Mirrors the config that `git xet install` writes, but scoped to this one
    clone so the broken interaction with `git clone --mirror` is avoided.
    """
    _run(["git", "-C", clone_dir, "config", "lfs.customtransfer.xet.path", "git-xet"])
    _run(["git", "-C", clone_dir, "config", "lfs.customtransfer.xet.args", "transfer"])
    _run(["git", "-C", clone_dir, "config", "lfs.customtransfer.xet.concurrent", "true"])


def _default_branch(url: str) -> str | None:
    """Return the source repo's default branch from its HEAD symref, if any."""
    out = _run_capture(["git", "ls-remote", "--symref", url, "HEAD"])
    for line in out.splitlines():
        if line.startswith("ref: refs/heads/") and line.endswith("\tHEAD"):
            return line.removeprefix("ref: refs/heads/").removesuffix("\tHEAD")
    return None


def copy_repo_with_history(src: str, dst: str, repo_type: str) -> int:
    """Clone the source branches/tags and push them to the destination.

    Preserves the full git history (commits, branches, tags) and LFS / xet
    objects. The destination repo must already exist (pre-created by the
    calling Space).

    Auth: the canonical HF pattern is HTTP Basic with username `hf_user` and
    the access token as the password. Rather than embedding the token in the
    clone/push URLs (which exposes it in the process command line, the temp
    `.git/config`, and any git output echoed into the Job logs), we register a
    git credential helper that reads the token from `HF_TOKEN` in the
    environment at request time. The token therefore never appears on a command
    line or in a config file; only the helper script (which references
    `$HF_TOKEN`) is stored. The Job container is single-tenant and ephemeral,
    so a global credential helper is safe here.
    """
    if repo_type not in REPO_TYPE_URL_PREFIX:
        log(f"unsupported type: {repo_type}")
        return 1

    token = os.environ.get("HF_TOKEN")
    if not token:
        log("HF_TOKEN is not set in the Job environment")
        return 1

    prefix = REPO_TYPE_URL_PREFIX[repo_type]
    base = HF_ENDPOINT.replace("https://", "")
    src_url = f"https://{base}/{prefix}{src}"
    dst_url = f"https://{base}/{prefix}{dst}"
    src_label = f"{prefix}{src}"
    dst_label = f"{prefix}{dst}"

    # Authenticate via a credential helper that reads HF_TOKEN from the env at
    # request time, so the token never lands in a remote URL, the process
    # command line, or .git/config. Applies to ls-remote, fetch, lfs and push.
    _run([
        "git",
        "config",
        "--global",
        "credential.helper",
        '!f() { echo "username=hf_user"; echo "password=${HF_TOKEN}"; }; f',
    ])

    _ensure_git_lfs()
    xet_available = _ensure_git_xet_binary()

    with tempfile.TemporaryDirectory() as tmp:
        clone_dir = os.path.join(tmp, "repo.git")

        default_branch = _default_branch(src_url)

        # Avoid `git clone --mirror`: on HF it also fetches internal `refs/pr/*`
        # refs, and some repos fail during smart-HTTP pack negotiation with
        # `fatal: expected 'packfile'` / `remote end hung up unexpectedly`.
        # Normal repo duplication only needs branches + tags.
        log(f"git init --bare {src_label} ...")
        _run(["git", "init", "--bare", clone_dir])
        _run(["git", "remote", "add", "origin", src_url], cwd=clone_dir)
        if default_branch:
            _run(["git", "symbolic-ref", "HEAD", f"refs/heads/{default_branch}"], cwd=clone_dir)

        log("git fetch origin (branches + tags only)")
        _run(
            [
                "git",
                "fetch",
                "origin",
                "+refs/heads/*:refs/heads/*",
                "+refs/tags/*:refs/tags/*",
            ],
            cwd=clone_dir,
        )

        log("git lfs fetch --all (pull LFS blobs into the local mirror)")
        _run(["git", "lfs", "fetch", "--all"], cwd=clone_dir)

        # Now that the clone is safely complete, register xet as a local LFS
        # custom-transfer agent. The LFS push step will pick it up for
        # xet-supporting destinations and fall back to `basic` otherwise.
        if xet_available:
            log("enabling xet for push (local config only)")
            _enable_xet_for_push(clone_dir)

        log(f"adding destination remote {dst_label}")
        _run(["git", "remote", "add", "dst", dst_url], cwd=clone_dir)

        log("git lfs push --all dst (upload blobs before refs)")
        _run(["git", "lfs", "push", "--all", "dst"], cwd=clone_dir)

        log("git push dst (branches + tags, force-update semantics)")
        _run(
            [
                "git",
                "push",
                "dst",
                "+refs/heads/*:refs/heads/*",
                "+refs/tags/*:refs/tags/*",
            ],
            cwd=clone_dir,
        )

    log("mirror complete.")
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description="Duplicate a Hugging Face repo or bucket.")
    parser.add_argument(
        "--type",
        required=True,
        choices=["bucket", "model", "dataset", "space"],
        help="Repo type. Source and destination must be the same type.",
    )
    parser.add_argument("--src", required=True, help="Source repo/bucket id (e.g. `user/name`).")
    parser.add_argument("--dst", required=True, help="Destination repo/bucket id.")
    args = parser.parse_args()

    if args.type == "bucket":
        return copy_bucket()
    return copy_repo_with_history(args.src, args.dst, repo_type=args.type)


if __name__ == "__main__":
    sys.exit(main())