"""Disk-backed hop store: resume, single-shot re-roll, and lower peak RAM.

Hops are causally dependent -- hop N is rendered from hop N-1's tail -- so the
cache key **chains**: each hop's key mixes in the previous hop's key. Editing
shot 1 therefore invalidates 2..N automatically, which is correct and must be
surfaced in the UI, because otherwise it reads as a bug.

**Why 16-bit.** A cached hop's last frame becomes the next hop's Qwen pin and
its AddGuide guide. Round-tripping float32 through 8-bit would make a resumed
chain diverge from an uninterrupted one -- the cache would change the output,
which defeats the point. FFV1 at `rgb48le` keeps ~16 bits per channel, which is
far below the VAE's own noise floor, so a resumed hop is indistinguishable from
a fresh one. The cost is roughly 2x the bytes of an 8-bit lossless encode, and
FFV1 still compresses it well.

**On the RAM claim.** The node's IMAGE output is the whole clip, so the final
tensor is unavoidably full size. What the store removes is the *double and
triple buffering* during the loop: today `master_imgs` grows by concatenation
(which allocates a new full-size tensor every hop) while `prev_imgs` and `imgs`
are also live. Streaming to disk keeps one hop plus the overlap tail resident
and concatenates once at the end.
"""

import hashlib
import json
import os
import shutil
import subprocess
import time

import numpy as np
import torch

TAG = "HandTieClips"

VIDEO_EXT = ".mkv"
AUDIO_EXT = ".npy"
META_EXT = ".json"
LATENT_EXT = ".latent.pt"


def _ffmpeg():
    exe = shutil.which("ffmpeg")
    if not exe:
        raise RuntimeError(
            f"{TAG}: ffmpeg is not on PATH. The hop store needs it to write "
            f"lossless FFV1. Install ffmpeg or set cache to off."
        )
    return exe


def tensor_digest(t):
    """Cheap, order-sensitive digest of a tensor's actual bytes."""
    if t is None:
        return "none"
    a = t.detach().cpu().contiguous().numpy()
    h = hashlib.sha256()
    h.update(str(a.shape).encode())
    h.update(str(a.dtype).encode())
    h.update(a.tobytes())
    return h.hexdigest()[:16]


def audio_digest(a):
    """Digest an AUDIO input -- ``{"waveform": Tensor, "sample_rate": int}``.

    AUDIO is a dict, not a tensor, so `tensor_digest` cannot take it: `.detach`
    on a dict raises AttributeError, which is what wiring `voice` with the hop
    cache on used to do before hop 1 ever started. Sample rate is part of the
    identity -- the same waveform at a different rate is different audio.
    """
    if a is None:
        return "none"
    if not isinstance(a, dict):
        return tensor_digest(a)
    h = hashlib.sha256()
    h.update(tensor_digest(a.get("waveform")).encode())
    h.update(str(a.get("sample_rate")).encode())
    return h.hexdigest()[:16]


def hop_key(prev_key, payload):
    """Chained content key. `payload` must contain everything that changes pixels.

    Anything omitted here is something the cache will fail to notice, so err
    toward including it.
    """
    h = hashlib.sha256()
    h.update((prev_key or "root").encode())
    h.update(json.dumps(payload, sort_keys=True, default=str).encode())
    return h.hexdigest()[:24]


class HopStore:
    def __init__(self, root, budget_gb=20.0, fps=24):
        self.root = str(root)
        self.budget = float(budget_gb) * (1024 ** 3)
        self.fps = int(fps)
        os.makedirs(self.root, exist_ok=True)

    # -- paths ------------------------------------------------------------
    def _p(self, key, ext):
        return os.path.join(self.root, key + ext)

    def has(self, key):
        return all(os.path.exists(self._p(key, e))
                   for e in (VIDEO_EXT, AUDIO_EXT, META_EXT))

    # -- write ------------------------------------------------------------
    def put(self, key, imgs, wav, sr, meta=None, latent=None):
        """imgs: float [N,H,W,3] in 0..1 on cpu. wav: float [.., C, S].

        `latent` is this hop's sampler output and is optional; see the comment
        at the write below for why storing it is what makes the cache useful
        past hop 1.
        """
        n, hgt, wid = int(imgs.shape[0]), int(imgs.shape[1]), int(imgs.shape[2])
        vid_tmp = self._p(key, VIDEO_EXT + ".part")
        cmd = [
            _ffmpeg(), "-y", "-v", "error",
            "-f", "rawvideo", "-pix_fmt", "rgb48le",
            "-s", f"{wid}x{hgt}", "-r", str(self.fps),
            "-i", "-",
            "-c:v", "ffv1", "-level", "3", "-coder", "1", "-context", "1",
            "-pix_fmt", "rgb48le",
            # The .part suffix defeats extension-based format detection, so the
            # muxer is named explicitly. Writing to .part and renaming on success
            # keeps a killed render from leaving a half-file that `has()` trusts.
            "-f", "matroska", vid_tmp,
        ]
        proc = subprocess.Popen(cmd, stdin=subprocess.PIPE,
                                stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
        try:
            # Frame at a time: never materialise a second full-size copy.
            for i in range(n):
                f = (imgs[i].clamp(0, 1) * 65535.0).round().to(torch.int32)
                proc.stdin.write(f.numpy().astype("<u2").tobytes())
            proc.stdin.close()
        except BrokenPipeError:
            pass
        err = proc.stderr.read().decode(errors="replace")
        if proc.wait() != 0:
            raise RuntimeError(f"{TAG}: FFV1 encode failed for hop {key}: {err.strip()}")
        os.replace(vid_tmp, self._p(key, VIDEO_EXT))

        np.save(self._p(key, AUDIO_EXT),
                wav.detach().cpu().contiguous().numpy().astype(np.float32))

        # The sampler latent, not just the decoded frames. Without it a cache
        # hit leaves the caller's `prev_sampled` empty, so the *next* hop
        # predicts the AddGuide pixel fallback instead of Motion-Context --
        # which is part of its key, so that key stops matching what was stored.
        # The cache could therefore never hit past hop 1, and the hop after a
        # hit was joined by the inferior mechanism, making `cache_hops=on`
        # actively worse than off. Optional and best-effort: a hop whose latent
        # will not serialise is still worth caching for its frames, it just
        # cannot seed a latent join.
        if latent is not None:
            tmp = self._p(key, LATENT_EXT + ".part")
            try:
                torch.save(latent, tmp)
                os.replace(tmp, self._p(key, LATENT_EXT))
            except Exception as e:  # noqa: BLE001
                print(f"[{TAG}] hop {key[:8]}: latent not cached ({e!r}); a hit "
                      f"on this hop will fall back to the pixel pin", flush=True)
                try:
                    os.remove(tmp)
                except OSError:
                    pass

        info = dict(meta or {})
        info.update({"frames": n, "height": hgt, "width": wid,
                     "sample_rate": int(sr), "fps": self.fps,
                     "written": time.time()})
        with open(self._p(key, META_EXT), "w", encoding="utf-8") as fh:
            json.dump(info, fh, indent=1)
        return info

    # -- read -------------------------------------------------------------
    def _get_latent(self, key):
        """The stored sampler latent, or None when this entry has none.

        None is always safe: the caller falls back to the pixel pin, which is
        what every cache hit did before latents were stored. Entries written by
        an older build simply have no sidecar, so they keep working.
        """
        path = self._p(key, LATENT_EXT)
        if not os.path.exists(path):
            return None
        try:
            # weights_only=False: this is a latent *dict*, written by our own
            # put() into ComfyUI's temp directory, not a downloaded checkpoint.
            return torch.load(path, map_location="cpu", weights_only=False)
        except Exception as e:  # noqa: BLE001
            print(f"[{TAG}] hop {key[:8]}: cached latent unreadable ({e!r}); "
                  f"falling back to the pixel pin", flush=True)
            return None

    def get(self, key):
        """-> (imgs float32 [N,H,W,3], wav, sample_rate, latent|None) or None."""
        if not self.has(key):
            return None
        with open(self._p(key, META_EXT), encoding="utf-8") as fh:
            info = json.load(fh)
        n, hgt, wid = int(info["frames"]), int(info["height"]), int(info["width"])
        cmd = [
            _ffmpeg(), "-v", "error", "-i", self._p(key, VIDEO_EXT),
            "-f", "rawvideo", "-pix_fmt", "rgb48le", "-",
        ]
        want = n * hgt * wid * 3 * 2
        # communicate(), not sequential reads: draining stdout to EOF while
        # stderr is an unread pipe deadlocks the moment ffmpeg emits more than
        # the pipe buffer on stderr -- which is exactly what a corrupt FFV1
        # does, i.e. the one case where the error actually matters.
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        raw, err_raw = proc.communicate()
        err = err_raw.decode(errors="replace")
        if proc.returncode != 0:
            raise RuntimeError(f"{TAG}: FFV1 decode failed for hop {key}: {err.strip()}")
        if len(raw) != want:
            raise RuntimeError(
                f"{TAG}: cached hop {key} is {len(raw)} bytes, expected {want}. "
                f"Delete the cache entry and re-render.")
        arr = np.frombuffer(raw, dtype="<u2").reshape(n, hgt, wid, 3)
        imgs = torch.from_numpy(arr.astype(np.float32) / 65535.0)
        wav = torch.from_numpy(np.load(self._p(key, AUDIO_EXT)))
        os.utime(self._p(key, VIDEO_EXT), None)  # LRU touch
        return imgs, wav, int(info["sample_rate"]), self._get_latent(key)

    # -- shot locks -------------------------------------------------------
    # A locked shot reuses whatever it last rendered even when its inputs
    # change, so it needs a name that survives an edit -- the content key by
    # definition does not. The pointer maps a stable shot name to the key that
    # shot last produced.
    def _ptr(self, name):
        safe = "".join(c if c.isalnum() or c in "-_" else "_" for c in str(name))
        return os.path.join(self.root, "ptr_" + safe + ".json")

    def set_pointer(self, name, key):
        with open(self._ptr(name), "w", encoding="utf-8") as fh:
            json.dump({"key": key, "written": time.time()}, fh)

    def get_pointer(self, name):
        try:
            with open(self._ptr(name), encoding="utf-8") as fh:
                key = json.load(fh).get("key")
        except (OSError, ValueError):
            return None
        return key if key and self.has(key) else None

    # -- housekeeping -----------------------------------------------------
    def entries(self):
        out = []
        for fn in os.listdir(self.root):
            if not fn.endswith(VIDEO_EXT):
                continue
            key = fn[: -len(VIDEO_EXT)]
            try:
                st = os.stat(self._p(key, VIDEO_EXT))
            except OSError:
                continue
            size = st.st_size
            for e in (AUDIO_EXT, META_EXT, LATENT_EXT):
                try:
                    size += os.path.getsize(self._p(key, e))
                except OSError:
                    pass
            out.append((key, size, st.st_mtime))
        return out

    def sweep(self, keep=()):
        """Evict least-recently-used entries until under budget.

        `keep` names hops this run still needs, so a budget smaller than one
        chain cannot delete the hop that is about to be read back.
        """
        items = self.entries()
        total = sum(s for _, s, _ in items)
        if total <= self.budget:
            return 0
        keep = set(keep)
        freed = 0
        for key, size, _ in sorted(items, key=lambda r: r[2]):
            if total - freed <= self.budget:
                break
            if key in keep:
                continue
            for e in (VIDEO_EXT, AUDIO_EXT, META_EXT, LATENT_EXT):
                try:
                    os.remove(self._p(key, e))
                except OSError:
                    pass
            freed += size
        if freed:
            print(f"[{TAG}] cache: evicted {freed / 1024 ** 3:.2f} GB "
                  f"(budget {self.budget / 1024 ** 3:.1f} GB)", flush=True)
        return freed