"""Bucket persistence for the zero-GPU training Space.

Only *weights and optimizer/trainer state* are persisted to an Hugging Face
Storage Bucket. Training data (the uploaded "trace") is never stored — it is
parsed into an in-memory ``datasets`` object and discarded after the run.

The bucket layout for one experiment (one model + one method) is::

    <checkpoint_id>/
        config.json, tokenizer.json, ...      # model + LoRA adapter weights
        trainer_state.json                     # global step, LR schedule, log
        optimizer.pt                           # Adam / AdamW moments  <-- the key
        scheduler.pt                           # LR scheduler state
        rng.pt                                 # RNG seed state
        scheduler.pt / trainer_state.json ...  # everything HF Trainer saves

``trainer_state.json`` records ``global_step`` so the next call knows exactly
where to resume. ``state.json`` is a tiny index of the most recent checkpoint
so the UI can show progress and resume without scanning the bucket.
"""

from __future__ import annotations

import json
import os
import tempfile
from pathlib import Path

from huggingface_hub import (
    batch_bucket_files,
    create_bucket,
    download_bucket_files,
    list_bucket_tree,
)

# Files a HF Trainer checkpoint that is safe to persist as "weights + state".
# Everything the trainer writes into <output_dir>/checkpoint-<step> is here.
PERSIST_GLOB = (
    "config.json",
    "tokenizer.json",
    "tokenizer_config.json",
    "special_tokens_map.json",
    "tokenizer.model",
    "generation_config.json",
    "adapter_config.json",
    "adapter_model.safetensors",
    "adapter_model.bin",
    "model.safetensors",
    "model.safetensors.index.json",
    "pytorch_model.bin",
    "trainer_state.json",
    "optimizer.pt",
    "scheduler.pt",
    "rng.pt",
    "training_args.bin",
)


def _token() -> str | None:
    return os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")


def ensure_bucket(bucket_id: str) -> None:
    """Create the bucket if it does not exist. Namespace-qualified id is fine
    (``username/my-bucket``); a bare name is created in the current user's
    namespace."""
    try:
        create_bucket(bucket_id=bucket_id, exist_ok=True, token=_token())
    except Exception:
        # Already exists / permission nuance — the real check is a later upload.
        pass


def checkpoint_paths(bucket_id: str) -> list[str]:
    """Return the sorted list of ``<checkpoint_id>`` prefixes currently in the
    bucket (i.e. every completed, persisted step)."""
    paths: list[str] = set()
    try:
        for entry in list_bucket_tree(bucket_id=bucket_id, prefix="", token=_token()):
            p = getattr(entry, "path", "") or ""
            # paths look like ``checkpoint-1/adapter_model.safetensors``; the
            # checkpoint id is the first path segment.
            first = p.split("/", 1)[0]
            if first.startswith("checkpoint-"):
                paths.add(first)
            elif p == "checkpoint-1" or first == "checkpoint-1":  # edge: bare
                paths.add(p)
    except Exception:
        pass
    return sorted(paths, key=_step_of)


def _step_of(prefix: str) -> int:
    try:
        return int(prefix.split("-")[-1])
    except Exception:
        return -1


def list_latest_checkpoint(bucket_id: str) -> str | None:
    paths = checkpoint_paths(bucket_id)
    return paths[-1] if paths else None


def push_checkpoint(bucket_id: str, local_dir: Path, checkpoint_id: str) -> bool:
    """Upload every weight/state file under ``local_dir`` into
    ``<bucket_id>/<checkpoint_id>/``. Returns True on success."""
    add: list[tuple[Path, str]] = []
    for path in local_dir.rglob("*"):
        if path.is_file():
            dest = f"{checkpoint_id}/{path.relative_to(local_dir).as_posix()}"
            add.append((path, dest))
    if not add:
        return False
    try:
        batch_bucket_files(bucket_id, add=add, token=_token())
        return True
    except Exception:
        return False


def pull_checkpoint(bucket_id: str, checkpoint_id: str, local_dir: Path) -> bool:
    """Download a persisted checkpoint from the bucket back into ``local_dir``.

    The *source* paths passed to ``download_bucket_files`` must be the full
    bucket-relative paths (``<checkpoint_id>/<relative>``); the local *dest*
    strips the prefix. Returns True if at least one file was retrieved.

    Xet-backed buckets can have read-after-write latency, so the whole
    list→download→verify cycle is retried a few times before giving up.
    """
    import time

    for _attempt in range(8):
        try:
            entries = list(
                list_bucket_tree(bucket_id=bucket_id, prefix=f"{checkpoint_id}/", token=_token())
            )
        except Exception:
            entries = []
        files: list[tuple[str, str]] = []
        for e in entries:
            full = getattr(e, "path", "") or ""
            if full.startswith(f"{checkpoint_id}/"):
                rel = full[len(f"{checkpoint_id}/"):]
                files.append((full, f"{local_dir}/{rel}"))
        if files:
            try:
                # download needs the destination parent dirs to exist
                local_dir.mkdir(parents=True, exist_ok=True)
                for _src, dest in files:
                    Path(dest).parent.mkdir(parents=True, exist_ok=True)
                download_bucket_files(bucket_id, files=files, token=_token())
            except Exception:
                files = []
            # verify at least one file actually landed on disk
            for _src, dest in files:
                if Path(dest).exists():
                    return True
        # no file landed yet — Xet may still be settling; retry
        time.sleep(1.5)
    return False


def save_state_index(bucket_id: str, state: dict) -> None:
    """Persist a tiny JSON index of the run so the UI can resume without a full
    tree scan. ``state`` e.g. ``{"global_step": 12, "loss": 0.42, "method": "dpo"}``."""
    content = json.dumps(state, indent=2).encode()
    batch_bucket_files(bucket_id, add=[(content, "state.json")], token=_token())


def load_state_index(bucket_id: str) -> dict:
    local = Path(tempfile.gettempdir()) / "zero_gpu_state.json"
    try:
        download_bucket_files(bucket_id, files=[("state.json", str(local))], token=_token())
    except Exception:
        return {}
    try:
        return json.loads(local.read_text())
    except Exception:
        return {}


def save_job(bucket_id: str, job) -> None:
    """Persist a queued job to the bucket (opt-in queue persistence).

    Writes ``queue/<job_id>.json`` so a queued request survives a Space
    restart. ``job`` is a ``queue.Job`` (duck-typed for ``job_id`` / ``method``
    / ``rows``). Best-effort: never raises into the caller.
    """
    try:
        content = json.dumps(
            {
                "job_id": getattr(job, "job_id", ""),
                "method": getattr(job, "method", ""),
                "rows": getattr(job, "rows", []),
                "submitted_at": getattr(job, "submitted_at", 0.0),
            },
            indent=2,
        ).encode()
        batch_bucket_files(
            bucket_id,
            add=[(content, f"queue/{getattr(job, 'job_id', 'job')}.json")],
            token=_token(),
        )
    except Exception:
        pass


def delete_job(bucket_id: str, job_id: str) -> None:
    """Remove a completed job from the bucket queue (opt-in). Best-effort."""
    try:
        from huggingface_hub import delete_bucket_file

        delete_bucket_file(bucket_id, f"queue/{job_id}.json", token=_token())
    except Exception:
        pass