"""The artifact store: bytes a node produced, addressed by their content.

A typed message carries JSON, which is what lets the same value pass through
Redis, the work queue and the worker protocol unchanged. A model checkpoint is
not that, so it does not travel as a message — it is written here and the
message carries a reference to it.

Addressed by digest rather than by run, for three reasons. A sweep whose fifty
configs share one preprocessed input stores it once. A reference stays valid
however it is passed around, because it names content instead of a location.
And the digest is what a stage cache will compare, so building it in now is
what keeps that from being a change to the message contract later.
"""

from __future__ import annotations

import contextlib
import hashlib
import logging
import os
import tempfile
import time
from collections.abc import Iterable, Iterator
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

#: How much is read at a time when hashing or serving.
CHUNK = 1024 * 1024
DIGEST_PREFIX = "sha256:"


def is_reference(value: Any) -> bool:
    """Whether a message payload is an artifact reference."""
    return isinstance(value, dict) and str(value.get("digest", "")).startswith(
        DIGEST_PREFIX
    )


def valid_digest(digest: str) -> bool:
    """Guard for anything that reaches the filesystem from outside."""
    if not digest.startswith(DIGEST_PREFIX):
        return False
    body = digest[len(DIGEST_PREFIX) :]
    return len(body) == 64 and all(c in "0123456789abcdef" for c in body)


class ArtifactStore:
    """Content-addressed files under one directory."""

    def __init__(self, root: Path) -> None:
        self.root = root
        self.root.mkdir(parents=True, exist_ok=True)
        #: A ring of frames held in memory, for media a flow only shows live.
        #: Set by the engine at startup; every lookup below falls through to
        #: it, which is what lets a volatile reference be an ordinary one.
        self.volatile: VolatileStore | None = None

    def _path(self, digest: str) -> Path:
        body = digest[len(DIGEST_PREFIX) :]
        # Two levels, so a directory listing stays usable at a hundred thousand
        # artifacts.
        return self.root / body[:2] / body

    def put(
        self,
        chunks: Iterable[bytes],
        name: str = "",
        media_type: str = "",
        volatile: bool = False,
    ) -> dict[str, Any]:
        """Store a stream and return the reference to it.

        Written to a temporary file first and moved into place once the digest
        is known, so a half-written artifact never has a name anyone can find.
        A file already there is left alone: identical content is identical.

        ``volatile`` puts it in the ring instead, where it is held in memory
        and falls out once newer frames need the room. The reference is the
        same shape either way — what differs is how long the bytes last, and a
        caller that wants one kept returns it from a run.
        """
        if volatile and self.volatile is not None:
            return self.volatile.put(chunks, name=name, media_type=media_type)
        digester = hashlib.sha256()
        size = 0
        handle = tempfile.NamedTemporaryFile(dir=self.root, delete=False)
        try:
            with handle:
                for chunk in chunks:
                    digester.update(chunk)
                    size += len(chunk)
                    handle.write(chunk)
            digest = DIGEST_PREFIX + digester.hexdigest()
            target = self._path(digest)
            target.parent.mkdir(parents=True, exist_ok=True)
            if target.exists():
                os.unlink(handle.name)
            else:
                # Same content from two nodes at once is one rename winning and
                # the other replacing a byte-identical file.
                os.replace(handle.name, target)
        except BaseException:
            with contextlib.suppress(OSError):
                os.unlink(handle.name)
            raise
        return {
            "digest": digest,
            "size": size,
            "media_type": media_type or "application/octet-stream",
            "name": name,
        }

    def put_file(
        self,
        path: Path,
        media_type: str = "",
        name: str = "",
        volatile: bool = False,
    ) -> dict[str, Any]:
        with path.open("rb") as handle:
            return self.put(
                iter(lambda: handle.read(CHUNK), b""),
                name=name or path.name,
                media_type=media_type,
                volatile=volatile,
            )

    def path(self, digest: str) -> Path | None:
        """Where the bytes are, or None if this store does not have them.

        The ring is looked in second, so everything that resolves a digest —
        serving one over HTTP, checking a run input still exists, a panel's
        scope — reaches a live frame without knowing there are two stores.
        """
        if not valid_digest(digest):
            return None
        target = self._path(digest)
        if target.exists():
            return target
        if self.volatile is not None:
            return self.volatile.path(digest)
        return None

    def adopt(self, digest: str) -> bool:
        """Copy a volatile artifact into this store, so it outlives the ring.

        What makes "emitted media is not kept, returned media is" true: a run
        recording a reference calls this, and the frame stops being one the
        next few seconds can evict.
        """
        if self.volatile is None or not valid_digest(digest):
            return False
        if self._path(digest).exists():
            return True
        source = self.volatile.path(digest)
        if source is None:
            return False
        self.put_file(source)
        return True

    def read(self, digest: str) -> Iterator[bytes]:
        target = self.path(digest)
        if target is None:
            raise FileNotFoundError(digest)
        with target.open("rb") as handle:
            while chunk := handle.read(CHUNK):
                yield chunk

    def collect(self, keep: set[str], grace_s: float = 0.0) -> int:
        """Delete what nothing refers to any more. Returns how many went.

        The caller passes every digest still recorded; anything else in the
        store was produced by a run that has since been pruned, or never got a
        row at all because the run failed between writing and recording.

        ``grace_s`` spares anything written that recently. Storing bytes and
        recording the reference to them are two steps, and a sweep landing
        between them would take an artifact its run is about to name — so
        recent files are left for the next pass, by which time they are either
        referenced or genuinely orphaned.
        """
        removed = 0
        cutoff = time.time() - grace_s
        for entry in self.root.glob("*/*"):
            if not entry.is_file():
                continue
            if DIGEST_PREFIX + entry.name in keep:
                continue
            try:
                if grace_s > 0 and entry.stat().st_mtime > cutoff:
                    continue
                entry.unlink()
                removed += 1
            except OSError:
                logger.warning("Could not remove artifact %s", entry.name)
        return removed


class VolatileStore(ArtifactStore):
    """A bounded ring of artifacts, held wherever memory is cheaper than disk.

    A camera publishing ten frames a second is ten files a second, and on the
    wall panel this is built for that disk is an SD card. So the frames go to a
    memory-backed directory instead and the oldest fall out once the ring is
    full: nothing sweeps it, because a frame nobody kept is not worth a pass
    over the store to find.

    It is an ``ArtifactStore``, digest layout and all, which is what lets a
    frame be an ordinary reference — the dtype check, the panel's scope, a
    node opening one with ``load_artifact`` and the widget fetching one all
    work on it unchanged.
    """

    def __init__(self, root: Path, limit_bytes: int) -> None:
        super().__init__(root)
        self.limit_bytes = limit_bytes

    def put(
        self,
        chunks: Iterable[bytes],
        name: str = "",
        media_type: str = "",
        volatile: bool = False,
    ) -> dict[str, Any]:
        reference = super().put(chunks, name=name, media_type=media_type)
        self.trim()
        return reference

    def trim(self) -> int:
        """Drop the oldest until the ring is inside its bound.

        Also called on a timer, because a worker in this container writes here
        itself and the engine never sees that ``put``.

        # ponytail: a scandir per trim. An in-memory index if a ring of
        # thousands of frames ever shows up in a profile.
        """
        if self.limit_bytes <= 0:
            return 0
        entries: list[tuple[float, int, Path]] = []
        total = 0
        for entry in self.root.glob("*/*"):
            try:
                stat = entry.stat()
            except OSError:
                continue
            if not entry.is_file():
                continue
            entries.append((stat.st_mtime, stat.st_size, entry))
            total += stat.st_size
        if total <= self.limit_bytes:
            return 0

        removed = 0
        for _mtime, size, entry in sorted(entries):
            if total <= self.limit_bytes:
                break
            try:
                entry.unlink()
            except OSError:
                continue
            total -= size
            removed += 1
        return removed