# Hub sync for the mix build: the instance is scratch, so nothing that took an hour to produce may live
# only there (§3.13 "never let the instance be the only copy", applied to a CPU job rather than a
# checkpoint).
#
# Granularity is one *source*: staging a source takes 5-40 minutes, so publishing its shards the moment
# it finishes bounds the loss from any interruption to the one source in flight, and a resume can skip
# every source already on the Hub without re-reading it.
#
# The stage repo is private-in-practice but public by project rule; it holds no credentials and no
# benchmark material -- only token ids from the approved mix.

import json
import os


def repo_id_from_env(default=None):
    return os.environ.get("OUNCE100M_STAGE_REPO") or default


def api():
    import os

    import ounce100m_credentials

    ounce100m_credentials.install()
    from huggingface_hub import HfApi

    return HfApi(token=os.environ["HF_TOKEN"])


def ensure_repo(repo):
    # Must install() first: the module-level create_repo() resolves its token from the environment, and
    # this is the FIRST hub call a build makes -- before any api() has run. Calling it unauthenticated
    # is what produced the misleading 401 in the cold-resume rehearsal (E-017), and the diagnosis in
    # these notes was wrong about the cause until p2-token-scope proved dataset writes are permitted.
    import os

    api()
    from huggingface_hub import create_repo

    # hub 1.11 (the image's version) takes the owner inside repo_id and has no `namespace` kwarg (E-013)
    create_repo(repo_id=repo, repo_type="dataset", private=False, exist_ok=True,
                token=os.environ["HF_TOKEN"])
    return repo


class StageListingError(RuntimeError):
    """Raised when the stage repo's contents cannot be determined.

    Returning an empty set here would be the single most dangerous shortcut in this pipeline: an empty
    listing is indistinguishable from "nothing staged yet", so a *broken* listing makes a resume
    silently re-stage everything (wasting hours) or, worse, makes a merge believe there is nothing to
    merge while still looking consistent. Callers must see the difference between empty and unknown."""


def remote_stage_files(repo):
    """Relative paths under stage/ that exist on the Hub. Raises rather than returning {} on failure."""
    import json
    import urllib.request

    errors = []
    a = api()
    try:
        files = a.list_repo_files(repo_id=repo, repo_type="dataset", revision="main")
        return {f for f in files if f.startswith("stage/")}
    except Exception as e:
        errors.append(f"list_repo_files: {type(e).__name__}: {str(e)[:160]}")
        # A freshly created repo has no commits, so "main" does not resolve and list_repo_files raises
        # for a reason that genuinely means empty. Distinguish it by asking whether the repo exists at
        # all -- NOT by swallowing the exception, which is what made the first version ambiguous.
        try:
            info = a.repo_info(repo_id=repo, repo_type="dataset")
            if not getattr(info, "sha", None):
                return set()
        except Exception as e2:
            errors.append(f"repo_info: {type(e2).__name__}: {str(e2)[:120]}")

    # REST fallback. Note the path goes in the query string, not the URL body -- the first version put
    # it in the path, got a 404, and the bare `except` turned that into a silent empty listing (E-018).
    for url in (f"https://huggingface.co/api/datasets/{repo}/tree/main?recursive=true",
                f"https://huggingface.co/api/datasets/{repo}/tree/main"):
        try:
            with urllib.request.urlopen(url, timeout=90) as r:
                entries = json.loads(r.read().decode())
            paths = set()
            stack = [(e, "") for e in entries]
            while stack:
                e, prefix = stack.pop()
                p = f"{prefix}/{e['path']}".lstrip("/") if prefix else e.get("path", "")
                if e.get("type") == "dir":
                    stack.extend((c, p) for c in (e.get("children") or []))
                elif p.startswith("stage/"):
                    paths.add(p)
            return paths
        except Exception as e:
            errors.append(f"tree {url.split('/api/')[1]}: {type(e).__name__}: {str(e)[:140]}")

    raise StageListingError(f"cannot list {repo}: " + " | ".join(errors))


def upload_source(root, key, repo):
    """Publish stage/<key>/* in one commit. Returns the file names pushed."""
    a = api()
    local = os.path.join(root, "stage", key)
    files = sorted(f for f in os.listdir(local) if f.endswith((".bin", ".json")))
    if not files:
        return []
    a.upload_folder(folder_path=local, path_in_repo=f"stage/{key}", repo_id=repo,
                    repo_type="dataset", commit_message=f"stage {key}: {len(files)} shard(s)",
                    allow_patterns=["*.bin", "*.json"])
    return files


def download_missing(root, keys_wanted, repo, remote=None):
    """Fetch stage/<key>/*.bin for every wanted key not already on disk. Returns {key: [filenames]}."""
    from huggingface_hub import hf_hub_download

    remote = remote if remote is not None else remote_stage_files(repo)
    got = {}
    for key in keys_wanted:
        # .json must come down too: record.json is what tells the restore the shard list, token count
        # and drop stats. A .bin-only filter made every restored source look unknown and was the actual
        # cause of "restored 0 staged source(s)" in rehearsal v3, with the files plainly present.
        want = sorted(p for p in remote
                      if p.startswith(f"stage/{key}/") and p.endswith((".bin", ".json")))
        if not want:
            got[key] = []
            continue
        local_dir = os.path.join(root, "stage", key)
        os.makedirs(local_dir, exist_ok=True)
        have = {f"stage/{key}/{f}" for f in os.listdir(local_dir)}
        for rp in want:
            if rp in have:
                continue
            hf_hub_download(repo_id=repo, repo_type="dataset", filename=rp,
                            local_dir=root, force_download=False)
        got[key] = sorted(os.listdir(local_dir))
    return got


def status(repo):
    import urllib.request

    url = f"https://huggingface.co/api/datasets/{repo}"
    try:
        with urllib.request.urlopen(url, timeout=60) as r:
            d = json.loads(r.read().decode())
        return {"exists": True, "gated": d.get("gated"), "private": d.get("private"),
                "lastModified": d.get("lastModified"), "id": d.get("id")}
    except Exception as e:
        return {"exists": False, "error": f"{type(e).__name__}: {str(e)[:120]}"}