# Copyright © The Debusine Developers
# See the AUTHORS file at the top-level directory of this distribution
#
# This file is part of Debusine. It is subject to the license terms
# in the LICENSE file found in the top-level directory of this
# distribution. No part of Debusine, including this file, may be copied,
# modified, propagated, or distributed except according to the terms
# contained in the LICENSE file.

"""Image handling for executor backends."""

import datetime as dt
import logging
from pathlib import Path
from shutil import rmtree

from fasteners import InterProcessLock  # type: ignore[import-untyped]

from debusine.client.debusine import Debusine
from debusine.client.models import ArtifactResponse
from debusine.tasks.executors.base import ExecutorImageCategory, _backends
from debusine.tasks.models import (
    BackendType,
    ImageCacheUsageLog,
    ImageCacheUsageLogEntry,
)

CACHE_SIZE = 10
USAGE_LOG = "usage.json"

log = logging.getLogger(__name__)
# Alias datetime.now to make it easier to mock
now = dt.datetime.now


class ImageCache:
    """
    Download image artifacts to a local cache.

    Processed versions of an image may be stored for a backend.
    """

    image_cache_path: Path = Path.home() / "system-images"
    image_category: ExecutorImageCategory
    backend: BackendType | None = None
    workspace: str

    def __init__(
        self,
        debusine_api: Debusine,
        image_category: ExecutorImageCategory,
        workspace: str,
    ) -> None:
        """
        Instantiate an ImageCache.

        :param debusine_api: object to use the debusine client API
        :param image_category: artifact category of supported images
        """
        self.debusine_api = debusine_api
        self.image_category = image_category
        self.workspace = workspace

    def cache_artifact_path(self, artifact_id: int) -> Path:
        """Return the path to the image cache pool for artifact_id."""
        return self.image_cache_path / str(artifact_id)

    def cache_artifact_image_path(
        self,
        artifact_id: int,
        filename: str,
        backend: BackendType | None = None,
    ) -> Path:
        """Return the path to the cached image for filename, for backend."""
        cache_path = self.cache_artifact_path(artifact_id)
        if backend:
            cache_path = cache_path / "backends" / backend
        return cache_path / filename

    def image_artifact(self, artifact_id: int) -> ArtifactResponse:
        """Retrieve an Artifact."""
        r = self.debusine_api.artifact_get(
            artifact_id, workspace=self.workspace
        )
        if r.category != self.image_category.value:
            raise ValueError(
                f"Unexpected artifact type {r.category}; expected "
                f"{self.image_category}"
            )
        return r

    def _image_lock(self, artifact_id: int) -> InterProcessLock:
        """Context Manager to hold a RW lock for working on a cached image."""
        cache_path = self.cache_artifact_path(artifact_id)
        cache_path.mkdir(exist_ok=True)
        return InterProcessLock(cache_path / ".lock")

    def process_downloaded_image(self, artifact: ArtifactResponse) -> None:
        """
        Handle any necessary post-download processing of the image.

        Locate the downloaded image using cache_artifact_image_path() and store
        it for a backend, if the result needs to be stored in the cache.

        This only needs to be done once, after the image is first downloaded.
        """

    def image_in_cache(self, artifact: ArtifactResponse) -> bool:
        """Return whether or not artifact is available in the local cache."""
        return self.cache_artifact_image_path(
            artifact.id, artifact.data["filename"]
        ).exists()

    def _download_image_artifact(
        self, artifact_id: int, filename: str, destination: Path
    ) -> None:
        """Atomically download filename from artifact to destination."""
        temp_path = destination.parent / (".download.tmp." + destination.name)
        temp_path.unlink(missing_ok=True)
        self.debusine_api.download_artifact_file(
            artifact_id,
            path_in_artifact=filename,
            destination=temp_path,
            workspace=self.workspace,
        )
        temp_path.rename(destination)

    def download_image(self, artifact: ArtifactResponse) -> Path:
        """
        Make the image available locally.

        Fetch the image from artifact storage, if it isn't already available,
        and make it available locally.

        Return a path to the image or name, as appropriate for the backend.
        """
        self.image_cache_path.mkdir(exist_ok=True)

        filename = artifact.data["filename"]

        # Protect our image-cache directory structure
        if filename.startswith(".") or filename in ("backends", USAGE_LOG):
            raise ValueError("Image has an invalid filename")

        path = self.cache_artifact_image_path(artifact.id, filename)

        with self._image_lock(artifact.id):
            if not self.image_in_cache(artifact):
                if not path.exists():
                    self._download_image_artifact(artifact.id, filename, path)
                self.process_downloaded_image(artifact)
            self.record_image_use(artifact)
        self.clean_up_old_images()
        return path

    def record_image_use(self, artifact: ArtifactResponse) -> None:
        """
        Record the use of an image.

        Executed while the image lock is held.
        """
        record_file = self.cache_artifact_image_path(artifact.id, USAGE_LOG)
        if record_file.exists():
            with record_file.open("r") as f:
                record = ImageCacheUsageLog.model_validate_json(f.read())
        else:
            record = ImageCacheUsageLog()

        if self.backend:
            record.backends.add(self.backend)

        record.usage = record.usage[-CACHE_SIZE + 1 :]
        record.usage.append(
            ImageCacheUsageLogEntry(
                filename=artifact.data["filename"],
                backend=self.backend,
                timestamp=now(dt.UTC),
            )
        )

        tmp_file = self.cache_artifact_image_path(
            artifact.id, f".tmp.{USAGE_LOG}"
        )
        with tmp_file.open("w") as f:
            f.write(record.model_dump_json())
        tmp_file.replace(record_file)

    def clean_up_old_images(self) -> None:
        """
        Purge old images from the cache.

        This implements simple LRU eviction, evicting all backend versions at
        once.
        """
        for artifact_id, backends in self._find_least_used_images():
            with self._image_lock(artifact_id):
                for backend in backends:
                    try:
                        backend_cls = _backends[backend]
                        if hasattr(backend_cls, "clean_up_image"):
                            backend_cls.clean_up_image(artifact_id)
                    except Exception:
                        # Clean-up is in a critical code path, don't fail, just
                        # log errors. If the log gets corrupted, we'll attempt
                        # a full clean-up including backends that may not exist.
                        log.warning(
                            "Failed to clean up %i backend %s",
                            artifact_id,
                            backend,
                            exc_info=True,
                        )
                rmtree(self.cache_artifact_path(artifact_id))

    def _find_least_used_images(self) -> list[tuple[int, set[BackendType]]]:
        """
        Find the least recently used images in the cache.

        Return a list of (artifact_id, [backends]) tuples.
        """
        last_used: dict[int, dt.datetime] = {}
        backends: dict[int, set[BackendType]] = {}
        unknown: list[int] = []

        for path in self.image_cache_path.iterdir():
            if not path.name.isdigit():
                continue
            artifact_id = int(path.name)
            with self._image_lock(artifact_id):
                record_file = path / USAGE_LOG
                if not record_file.exists():
                    log.warning("No usage log for artifact %i", artifact_id)
                    unknown.append(artifact_id)
                    continue
                try:
                    with record_file.open("r") as f:
                        record = ImageCacheUsageLog.model_validate_json(
                            f.read()
                        )
                except Exception:
                    log.warning(
                        "Failed to parse log for artifact %i", artifact_id
                    )
                    unknown.append(artifact_id)
                    continue
                last_used[artifact_id] = record.usage[-1].timestamp
                backends[artifact_id] = record.backends

        most_recently_used = sorted(
            last_used.items(), key=lambda item: item[1], reverse=True
        )

        least_recently_used = []

        for artifact_id, timestamp in most_recently_used[CACHE_SIZE:]:
            least_recently_used.append((artifact_id, backends[artifact_id]))

        for artifact_id in unknown:
            # Assume the worst
            least_recently_used.append((artifact_id, set(_backends.keys())))

        return least_recently_used
