"""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 threading
import time
from collections import OrderedDict
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)
        #: Frames held in memory, for media a flow only shows live. Set by the
        #: engine at startup unless the ring is turned off; 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. Without a ring it
        lands here like anything else.
        """
        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 stored bytes are, or None if this store does not have them.

        Only the durable store: a live frame has no file. ``has`` and ``read``
        are what reach both.
        """
        if not valid_digest(digest):
            return None
        target = self._path(digest)
        return target if target.exists() else None

    def has(self, digest: str) -> bool:
        """Whether the bytes are here, stored or held in the ring.

        Everything that only asks whether a digest resolves — a run input
        that has to still exist, a cached result — reaches a live frame
        without knowing there are two stores.
        """
        if self.path(digest) is not None:
            return True
        return self.volatile is not None and self.volatile.has(digest)

    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
        data = self.volatile.get(digest)
        if data is None:
            return False
        self.put([data])
        return True

    def delete(self, digest: str) -> bool:
        """Remove one artifact's bytes. Returns whether this store had them.

        The ring is not looked in: its frames fall out on their own.
        """
        if not valid_digest(digest):
            return False
        try:
            self._path(digest).unlink()
        except FileNotFoundError:
            return False
        return True

    def read(self, digest: str) -> Iterator[bytes]:
        target = self.path(digest)
        if target is None:
            data = self.volatile.get(digest) if self.volatile is not None else None
            if data is None:
                raise FileNotFoundError(digest)
            yield data
            return
        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:
    """Live frames, held in memory and bounded by size.

    A camera publishing five frames a second is five artifacts a second that
    nobody keeps, so they never reach a disk: each one is held here in the
    order it came and the oldest fall out once the bound is reached. Nothing
    sweeps it, and a ``put`` costs the same however full it is.

    ``root`` is a memory-backed directory beside it for the frames a worker in
    this container writes itself (``FLUKSIO_ARTIFACT_VOLATILE_DIR``), since a
    worker is another process. Read from there once and held here after, so
    a frame pushed to every open screen is read from the file once. The
    engine trims that directory on a timer.

    A frame is an ordinary reference all the same: the dtype check, the
    device's scope, a node opening one with ``load_artifact`` and the widget
    fetching one all work on it unchanged, through ``ArtifactStore``.
    """

    def __init__(self, root: Path, limit_bytes: int) -> None:
        self.root = root
        self.root.mkdir(parents=True, exist_ok=True)
        self.limit_bytes = limit_bytes
        self._frames: OrderedDict[str, bytes] = OrderedDict()
        self._size = 0
        self._lock = threading.Lock()

    def put(
        self,
        chunks: Iterable[bytes],
        name: str = "",
        media_type: str = "",
    ) -> dict[str, Any]:
        """Hold a frame and return the reference to it."""
        data = b"".join(chunks)
        digest = DIGEST_PREFIX + hashlib.sha256(data).hexdigest()
        self._hold(digest, data)
        return {
            "digest": digest,
            "size": len(data),
            "media_type": media_type or "application/octet-stream",
            "name": name,
        }

    def _hold(self, digest: str, data: bytes) -> None:
        with self._lock:
            if digest in self._frames:
                # The same bytes again are the newest frame again.
                self._frames.move_to_end(digest)
                return
            self._frames[digest] = data
            self._size += len(data)
            # The one just held stays, even past the bound: its reference is
            # about to be published.
            while self._size > self.limit_bytes and len(self._frames) > 1:
                _, dropped = self._frames.popitem(last=False)
                self._size -= len(dropped)

    def _file(self, digest: str) -> Path:
        body = digest[len(DIGEST_PREFIX) :]
        return self.root / body[:2] / body

    def get(self, digest: str) -> bytes | None:
        """The frame's bytes, or None once it has fallen out."""
        with self._lock:
            data = self._frames.get(digest)
        if data is not None or not valid_digest(digest):
            return data
        try:
            data = self._file(digest).read_bytes()
        except OSError:
            return None
        self._hold(digest, data)
        return data

    def has(self, digest: str) -> bool:
        with self._lock:
            if digest in self._frames:
                return True
        return valid_digest(digest) and self._file(digest).is_file()

    def spill(self, digest: str) -> None:
        """Write a frame held in memory to the directory, for a worker here.

        A node in this container opens a frame with ``load_artifact``, which
        looks in the directory: the worker is another process and cannot see
        this one's memory. Called before a call naming the frame. The
        directory's timer trims it like a worker's own frames.
        """
        with self._lock:
            data = self._frames.get(digest)
        if data is None:
            return
        target = self._file(digest)
        if target.exists():
            return
        target.parent.mkdir(parents=True, exist_ok=True)
        handle = tempfile.NamedTemporaryFile(dir=target.parent, delete=False)
        try:
            with handle:
                handle.write(data)
            os.replace(handle.name, target)
        except BaseException:
            with contextlib.suppress(OSError):
                os.unlink(handle.name)
            raise

    @property
    def size(self) -> int:
        """Bytes held in memory."""
        return self._size

    def trim(self) -> int:
        """Drop the oldest of the workers' frames until the directory is
        inside the bound. Returns how many went.

        A scan of the directory, so it runs on a timer and never per frame:
        only frames a worker in this container wrote are there.
        """
        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