"""Persist per-file index results across server restarts.

Entries are keyed by path and validated by the file's size and mtime. The
whole cache is tied to a fingerprint of the code that produced it: the
indexing sources of this server, the bundled tclint sources and the versions
of all bundled libraries. Any change to them, including a tclint update or a
local patch, discards the cache instead of loading stale results.
"""

from __future__ import annotations

import hashlib
import logging
import os
import pathlib
import pickle
import sys
import tempfile
import threading
import zlib
from typing import Any

LOGGER = logging.getLogger(__name__)

# Bump when the cached data layout changes without a source change above.
CACHE_FORMAT = 1
CACHE_FILE = "index-cache.pickle.z"

_SRC_DIR = pathlib.Path(__file__).resolve().parent.parent
_LIBS_DIR = _SRC_DIR.parent / "libs"

FileStat = tuple[int, int]


def code_fingerprint() -> str:
    digest = hashlib.sha256()
    digest.update(f"{CACHE_FORMAT}|{sys.version}".encode())
    sources = [
        _SRC_DIR / "lsp_tclserver.py",
        *sorted((_SRC_DIR / "tools").glob("*.py")),
        *sorted((_SRC_DIR / "plugins").glob("*.py")),
        *sorted((_LIBS_DIR / "tclint").rglob("*.py")),
    ]
    for source in sources:
        digest.update(source.relative_to(_SRC_DIR.parent).as_posix().encode())
        digest.update(source.read_bytes())
    for dist_info in sorted(_LIBS_DIR.glob("*.dist-info")):
        digest.update(dist_info.name.encode())
    return digest.hexdigest()


def file_stat(path: str) -> FileStat | None:
    try:
        stat = os.stat(path)
    except OSError:
        return None
    return stat.st_mtime_ns, stat.st_size


class IndexCache:
    def __init__(self, directory: pathlib.Path | None = None, fingerprint: str = ""):
        self._path = directory / CACHE_FILE if directory is not None else None
        self._fingerprint = fingerprint
        self._entries: dict[str, tuple[FileStat, Any]] = {}
        self._used: set[str] = set()
        self._dirty = False
        self._lock = threading.Lock()

    @classmethod
    def load(cls, directory: pathlib.Path | str | None) -> IndexCache:
        """Open the cache in `directory`; without one, nothing is persisted."""
        if not directory:
            return cls()
        cache = cls(pathlib.Path(directory), code_fingerprint())
        try:
            with open(cache._path, "rb") as file:
                fingerprint, entries = pickle.loads(zlib.decompress(file.read()))
        except FileNotFoundError:
            return cache
        except Exception as error:  # A damaged cache must never stop indexing.
            LOGGER.warning("Ignoring unreadable index cache %s: %s", cache._path, error)
            cache._dirty = True
            return cache
        if fingerprint == cache._fingerprint:
            cache._entries = entries
        else:
            cache._dirty = True
        return cache

    def get(self, path: str, stat: FileStat) -> Any | None:
        with self._lock:
            entry = self._entries.get(path)
            if entry is None or entry[0] != stat:
                return None
            self._used.add(path)
            return entry[1]

    def put(self, path: str, stat: FileStat, data: Any) -> None:
        if self._path is None:
            return
        with self._lock:
            self._entries[path] = (stat, data)
            self._used.add(path)
            self._dirty = True

    def save(self) -> None:
        """Write entries used in this session atomically; others are dropped."""
        if self._path is None:
            return
        with self._lock:
            if not self._dirty and self._used == self._entries.keys():
                return
            entries = {path: self._entries[path] for path in self._used if path in self._entries}
            self._entries = entries
            self._dirty = False
        try:
            self._path.parent.mkdir(parents=True, exist_ok=True)
            with tempfile.NamedTemporaryFile(dir=self._path.parent, delete=False) as file:
                data = pickle.dumps((self._fingerprint, entries), protocol=pickle.HIGHEST_PROTOCOL)
                # Pickled indexes are very repetitive; fast compression cuts ~90%.
                file.write(zlib.compress(data, 1))
            os.replace(file.name, self._path)
        except Exception as error:
            LOGGER.warning("Could not write index cache %s: %s", self._path, error)
            try:
                os.unlink(file.name)
            except (OSError, NameError):
                pass