"""Gradio Space UI that delegates the actual repo/bucket copy to an HF Job.

The Job is launched via `huggingface_hub.run_uv_job` using a user-provided
Hugging Face token. `worker.py` is uploaded with the Job (via
`LOCAL_FILES_ENCODED`). The Space submits, polls logs, and offers a cancel
button — no `hf-mount`, no Dockerfile.

Strategy by repo type:
  - bucket                   : two `Volume(type="bucket", ...)` mounts; worker
                               copies file by file.
  - model | dataset | space  : worker does `git clone --mirror` + `git push
                               --mirror` (with LFS and xet auto-installed) to
                               preserve full content + git history. The Job
                               flavor is picked at submission time to give the
                               container an ephemeral disk large enough to
                               hold the mirror (see `_pick_flavor`).

Only the content is migrated — no settings (privacy, gated, hardware, …) are
mirrored. The single exception is `space_sdk`, which is required by
`create_repo(repo_type="space")` and is therefore read from the source.
"""

from __future__ import annotations

import os

import gradio as gr
import trackio
from huggingface_hub import (
    Volume,
    bucket_info,
    cancel_job,
    create_bucket,
    create_repo,
    dataset_info,
    fetch_job_logs,
    inspect_job,
    list_jobs_hardware,
    model_info,
    run_uv_job,
    space_info,
    whoami,
)

JOB_FLAVOR = os.environ.get("JOB_FLAVOR", "cpu-upgrade")
JOB_TIMEOUT = os.environ.get("JOB_TIMEOUT", "2h")
DISK_OVERHEAD_FACTOR = float(os.environ.get("DISK_OVERHEAD_FACTOR", "2.0"))
TRACKIO_PROJECT = os.environ.get("TRACKIO_PROJECT", "hf-cross-region-copy")
TRACKIO_SPACE_ID = os.environ.get("TRACKIO_SPACE_ID") or None
TRACKIO_ENABLED = os.environ.get("TRACKIO_ENABLED", "1").lower() not in {
    "0",
    "false",
    "no",
}

REPO_TYPES = ("bucket", "model", "dataset", "space")
REGIONS = ("us", "eu")

WORKER_SCRIPT = "worker.py"
SRC_MOUNT = "/data/src"
DST_MOUNT = "/data/dst"


def _parse_storage_bytes(s: str) -> int:
    """Parse strings like '20 GB', '100 GB', '1 TB' into bytes. Returns 0 on
    unparseable input so the calling code can ignore the flavor."""
    parts = (s or "").strip().split()
    if len(parts) != 2:
        return 0
    try:
        value = float(parts[0])
    except ValueError:
        return 0
    multipliers = {"KB": 1e3, "MB": 1e6, "GB": 1e9, "TB": 1e12}
    return int(value * multipliers.get(parts[1].upper(), 0))


def _source_repo_used_storage(repo_type: str, source_id: str, *, token: str) -> int:
    """Bytes used by the source repo on the Hub, or 0 if unknown."""
    info_fn = {"model": model_info, "dataset": dataset_info, "space": space_info}[repo_type]
    used = info_fn(source_id, token=token).used_storage
    return int(used or 0)


def _source_used_storage(repo_type: str, source_id: str, *, token: str) -> int:
    if repo_type == "bucket":
        return int(bucket_info(source_id, token=token).size or 0)
    return _source_repo_used_storage(repo_type, source_id, token=token)


def _pick_flavor(repo_type: str, source_id: str, *, token: str) -> tuple[str, int, int]:
    """Pick a Job flavor.

    - bucket → `JOB_FLAVOR` (no disk scaling, the worker just streams between
      two Volume mounts).
    - others → among all flavors whose ephemeral disk is at least
      `used_storage × DISK_OVERHEAD_FACTOR`, pick the **cheapest per minute**
      (`unit_cost_usd`). CPU and GPU flavors compete on equal footing — in
      practice CPU usually wins because GPU flavors carry a GPU markup, but
      we don't assume it. Raises a `gr.Error` if no flavor fits.

    Returns `(flavor_name, needed_bytes, picked_disk_bytes)`.
    """
    if repo_type == "bucket":
        return JOB_FLAVOR, 0, 0

    used = _source_repo_used_storage(repo_type, source_id, token=token)
    needed = int(used * DISK_OVERHEAD_FACTOR)

    fits = [
        (f, _parse_storage_bytes(f.ephemeral_storage))
        for f in list_jobs_hardware(token=token)
    ]
    fits = [(f, d) for f, d in fits if d >= needed]
    if not fits:
        raise gr.Error(
            f"No Job flavor has enough ephemeral disk to mirror this repo: "
            f"need ≥ {needed / 1e9:.1f} GB (source = {used / 1e9:.1f} GB × "
            f"{DISK_OVERHEAD_FACTOR}× overhead)."
        )

    f, disk = min(fits, key=lambda fb: fb[0].unit_cost_usd)
    return f.name, needed, disk


def _ensure_destination(
    *,
    repo_type: str,
    source_id: str,
    repo_id: str,
    region: str,
    token: str,
) -> None:
    """Create the destination — name + region only. **No metadata migration.**

    For spaces, `space_sdk` is the one exception: the create_repo API rejects
    `repo_type="space"` without it, so we read it from the source. Everything
    else (private, gated, hardware, variables, …) defaults to whatever the
    Hub does for a brand-new repo, and the duplicated content (git history,
    LFS/xet objects, card_data inside README) lands via the worker's mirror push.
    """
    if repo_type == "bucket":
        create_bucket(
            bucket_id=repo_id,
            region=region,
            exist_ok=True,
            token=token,
        )
        return

    kwargs: dict = {}
    if repo_type == "space":
        sdk = space_info(source_id, token=token).sdk
        if not sdk:
            raise gr.Error(
                f"Could not determine the SDK of source space `{source_id}`."
            )
        kwargs["space_sdk"] = sdk

    create_repo(
        repo_id=repo_id,
        repo_type=repo_type,
        region=region,
        exist_ok=True,
        token=token,
        **kwargs,
    )


def _track_transfer_submission(
    *,
    repo_type: str,
    source: str,
    destination_region: str,
    job_id: str,
    token: str,
) -> None:
    if not TRACKIO_ENABLED:
        return

    try:
        volume_bytes = _source_used_storage(repo_type, source, token=token)
        destination_metric = f"to_{destination_region}"
        init_kwargs = {
            "project": TRACKIO_PROJECT,
            "config": {
                "job_id": job_id,
                "repo_type": repo_type,
                "destination_region": destination_region,
            },
        }
        if TRACKIO_SPACE_ID:
            init_kwargs["space_id"] = TRACKIO_SPACE_ID

        trackio.init(**init_kwargs)
        try:
            trackio.log(
                {
                    "transfer_submissions": 1,
                    f"transfer_count_{destination_metric}": 1,
                    "volume_bytes": volume_bytes,
                    f"volume_{destination_metric}_bytes": volume_bytes,
                }
            )
        finally:
            trackio.finish()
    except Exception:
        # Tracking should never block a duplication job from starting.
        return


def submit_transfer(
    token: str,
    repo_type: str,
    source: str,
    destination: str,
    region: str,
):
    token = (token or "").strip()
    if not token:
        raise gr.Error(
            "Please provide a Hugging Face token with write access before starting a transfer."
        )

    if repo_type not in REPO_TYPES:
        raise gr.Error(f"Repo type must be one of {REPO_TYPES}.")
    source = (source or "").strip()
    destination = (destination or "").strip()
    if not source or not destination:
        raise gr.Error("Source and destination IDs are required (e.g. `username/my-repo`).")
    if source == destination:
        raise gr.Error("Source and destination must be different.")
    if region not in REGIONS:
        raise gr.Error(f"Region must be one of {REGIONS}.")

    try:
        user = whoami(token=token)
    except Exception as e:
        raise gr.Error(f"Invalid Hugging Face token or insufficient permissions: {e}")

    try:
        _ensure_destination(
            repo_type=repo_type,
            source_id=source,
            repo_id=destination,
            region=region,
            token=token,
        )
    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f"Failed to create destination {repo_type} `{destination}`: {e}")

    try:
        flavor, needed_bytes, disk_bytes = _pick_flavor(
            repo_type, source, token=token
        )
    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f"Failed to pick a Job flavor: {e}")

    script_args = ["--type", repo_type, "--src", source, "--dst", destination]

    volumes = None
    if repo_type == "bucket":
        volumes = [
            Volume(type="bucket", source=source, mount_path=SRC_MOUNT, read_only=True),
            Volume(type="bucket", source=destination, mount_path=DST_MOUNT, read_only=False),
        ]

    try:
        job = run_uv_job(
            script=WORKER_SCRIPT,
            script_args=script_args,
            volumes=volumes,
            secrets={"HF_TOKEN": token},
            flavor=flavor,
            timeout=JOB_TIMEOUT,
            token=token,
        )
    except Exception as e:
        raise gr.Error(f"Failed to submit job: {e}")

    _track_transfer_submission(
        repo_type=repo_type,
        source=source,
        destination_region=region,
        job_id=job.id,
        token=token,
    )

    if repo_type == "bucket":
        flavor_md = f"`{flavor}`"
    else:
        flavor_md = (
            f"`{flavor}` ({disk_bytes / 1e9:.0f} GB disk for "
            f"{needed_bytes / 1e9:.1f} GB needed)"
        )

    status_md = (
        f"**Status**: submitted\n\n"
        f"Token owner: **@{user['name']}**\n\n"
        f"Destination region: `{region}`\n\n"
        f"Job: [`{job.id}`]({job.url}) · type=`{repo_type}` · "
        f"{flavor_md} (timeout `{JOB_TIMEOUT}`)"
    )
    return job.id, token, status_md, ""


def poll_job(job_id: str, token: str):
    if not job_id or not token:
        return gr.update(), gr.update()

    try:
        info = inspect_job(job_id=job_id, token=token)
        log_lines = list(fetch_job_logs(job_id=job_id, follow=False, token=token))
    except Exception as e:
        return gr.update(), f"**Status**: poll error — {e}"

    stage = info.status.stage
    badge = {
        "RUNNING": "🟢 RUNNING",
        "COMPLETED": "✅ COMPLETED",
        "CANCELED": "🟠 CANCELED",
        "ERROR": "🔴 ERROR",
        "DELETED": "⚫ DELETED",
    }.get(stage, stage)

    status_md = f"**Status**: {badge} — [`{info.id}`]({info.url})"
    if info.status.message:
        status_md += f"\n\n> {info.status.message}"

    logs = "\n".join(log_lines)
    return logs, status_md


def cancel_transfer(job_id: str, token: str):
    if not job_id:
        return job_id, "**Status**: nothing to cancel"
    if not token:
        raise gr.Error("Please provide the same Hugging Face token used to start the job.")
    try:
        cancel_job(job_id=job_id, token=token)
    except Exception as e:
        raise gr.Error(f"Failed to cancel job: {e}")
    return "", f"**Status**: cancel requested for `{job_id}`"


with gr.Blocks(title="hf-cross-region-copy") as demo:
    gr.Markdown(
        "# hf-cross-region-copy\n"
        "Copy a Hugging Face bucket, model, dataset or space. Provide a Hugging Face "
        "token with write access; the transfer runs as an "
        "[HF Job](https://huggingface.co/docs/huggingface_hub/guides/jobs) using that token."
    )
    gr.Markdown(
        "For model, dataset and Space repos, preserving git history requires Git-over-HTTPS "
        "push access. Use a fine-grained token with write access to the source/destination "
        "namespace (and permission to run Jobs)."
    )

    with gr.Group() as form_group:
        token_input = gr.Textbox(
            label="Hugging Face token",
            type="password",
            placeholder="hf_...",
            info="Used for repo creation, Job submission, polling/cancel, and git push inside the Job.",
        )
        type_input = gr.Dropdown(
            label="Repo type",
            choices=list(REPO_TYPES),
            value="bucket",
        )

        with gr.Row():
            src_input = gr.Textbox(
                label="Source",
                placeholder="name/source",
                scale=1,
            )
            dst_input = gr.Textbox(
                label="Destination",
                placeholder="name/destination",
                scale=1,
            )

        region_input = gr.Dropdown(
            label="Destination region for new destinations",
            info=(
                "Used only when the destination is created. If it already exists, "
                "its current region is kept."
            ),
            choices=list(REGIONS),
            value="eu",
        )

        with gr.Row():
            submit_btn = gr.Button("Start transfer", variant="primary", scale=1)
            cancel_btn = gr.Button("Cancel", variant="stop", scale=0)

        status_md = gr.Markdown(value="**Status**: idle")
        log_box = gr.Textbox(
            label="Job logs",
            interactive=False,
            lines=20,
            max_lines=200,
            elem_id="log-box",
        )

        job_id_state = gr.State("")
        token_state = gr.State("")

    submit_btn.click(
        fn=submit_transfer,
        inputs=[token_input, type_input, src_input, dst_input, region_input],
        outputs=[job_id_state, token_state, status_md, log_box],
    )
    cancel_btn.click(
        fn=cancel_transfer,
        inputs=[job_id_state, token_state],
        outputs=[job_id_state, status_md],
    )

    log_timer = gr.Timer(value=2)
    log_timer.tick(fn=poll_job, inputs=[job_id_state, token_state], outputs=[log_box, status_md])


demo.queue()
demo.launch()