"""Models are loaded when first asked for and kept in a bounded cache.

The Space has 16 GB of memory and the registry lists more models each
release, so start-up loads only the default model and the others are loaded
on first use. Loaded models are kept, least recently used first out, while
their estimated memory stays within a budget (``TAXONOVA_WEB_MODEL_BUDGET_MB``,
default ``DEFAULT_BUDGET_MB``). The estimate is the checkpoint file size times
``MEMORY_PER_FILE_BYTE``. The model being asked for is always loaded, even if
it alone exceeds the budget.

Measured 2026-11-04 for the 26.11 models (taxonova 26.11.0, torch 2.5.1 CPU,
2 threads, one model per process): resident memory grew by 263 to 640 MB per
loaded model (1.2 to 1.3 times the file size; v4 334 MB, v4v5 640 MB), and
loading peaked about 0.35 GB above that. A 2,000-read job (rRNA screen and
classification) took 25 to 58 s and raised the process peak to 0.8 to 1.2 GB.

Jobs run one at a time in one worker thread, so a model is never evicted
while it is in use; the lock only protects the bookkeeping against the page
and health requests that read it.
"""
from __future__ import annotations

import ctypes
import gc
import logging
import os
import threading
from collections import OrderedDict
from pathlib import Path
from typing import Callable, Dict, List, Optional

logger = logging.getLogger(__name__)

#: Estimated memory per byte of checkpoint file: the measured 1.2 to 1.3 once
#: loaded, with margin for the load peak and allocator overhead.
MEMORY_PER_FILE_BYTE = 2.0
#: About 4 GB resident at the measured ratio, which leaves most of the
#: Space's 16 GB for a 20,000-sequence job, the rRNA screen and the server.
DEFAULT_BUDGET_MB = 6000
BUDGET_ENV = "TAXONOVA_WEB_MODEL_BUDGET_MB"


def budget_mb() -> int:
    try:
        return max(0, int(os.environ.get(BUDGET_ENV, DEFAULT_BUDGET_MB)))
    except ValueError:
        return DEFAULT_BUDGET_MB


def estimate_mb(path) -> float:
    return Path(path).stat().st_size * MEMORY_PER_FILE_BYTE / 2 ** 20


def _release_memory() -> None:
    gc.collect()
    try:  # hand freed heap pages back to the system (glibc only)
        ctypes.CDLL("libc.so.6").malloc_trim(0)
    except (OSError, AttributeError):
        pass


class ModelCache:
    """Loaded models by id, least recently used first, within a memory budget."""

    def __init__(self, loader: Callable[[str], object], budget: Optional[float] = None,
                 estimate: Callable[[str], float] = estimate_mb):
        self._loader = loader
        self._estimate = estimate
        self.budget = budget_mb() if budget is None else budget
        self._items: "OrderedDict[str, object]" = OrderedDict()
        self._sizes: Dict[str, float] = {}
        self._lock = threading.Lock()
        self.loads = 0
        self.evictions = 0

    def loaded_ids(self) -> List[str]:
        with self._lock:
            return list(self._items)

    def used_mb(self) -> float:
        with self._lock:
            return sum(self._sizes.values())

    def is_loaded(self, model_id: str) -> bool:
        with self._lock:
            return model_id in self._items

    def put(self, model_id: str, loaded, size_mb: float = 0.0) -> None:
        """Insert an already loaded model (tests, or a model loaded elsewhere)."""
        with self._lock:
            self._items[model_id] = loaded
            self._items.move_to_end(model_id)
            self._sizes[model_id] = size_mb

    def get(self, model_id: str, path) -> object:
        """The loaded model, loading it (and evicting others) if needed."""
        with self._lock:
            if model_id in self._items:
                self._items.move_to_end(model_id)
                return self._items[model_id]
        need = self._estimate(str(path))
        evicted = []
        with self._lock:
            while self._items and sum(self._sizes.values()) + need > self.budget:
                old, _ = self._items.popitem(last=False)
                self._sizes.pop(old, None)
                evicted.append(old)
        if evicted:
            self.evictions += len(evicted)
            logger.info("model cache: unloaded %s to make room for %s", ", ".join(evicted), model_id)
            _release_memory()
        logger.info("loading %s from %s (about %.0f MB)", model_id, path, need)
        loaded = self._loader(str(path))
        self.loads += 1
        with self._lock:
            self._items[model_id] = loaded
            self._sizes[model_id] = need
        return loaded