"""LoRA cache quota management.

Two helpers:
- `_lora_cache_usage_bytes()` — best-effort total size of files in the cache dir.
- `_evict_lora_cache_for_quota(needed)` — LRU-evict cache files (skipping any
  referenced by an in-memory LoRA store) until `needed` more bytes fit, or
  raise `gr.Error` if the quota still cannot accommodate the request.
"""
from __future__ import annotations

import os
import re

import gradio as gr

from state import (
    _CUSTOM_LoRA,
    _H3_MAX_LORA_CACHE_BYTES,
    _LORA_CACHE_DIR,
    _log,
)


def _lora_cache_usage_bytes() -> int:
    """Sum sizes of files under the LoRA cache dir (best-effort)."""
    total = 0
    try:
        if not os.path.isdir(_LORA_CACHE_DIR):
            return 0
        for name in os.listdir(_LORA_CACHE_DIR):
            p = os.path.join(_LORA_CACHE_DIR, name)
            if os.path.isfile(p) and not name.endswith(".tmp"):
                try:
                    total += os.path.getsize(p)
                except OSError:
                    pass
    except OSError:
        pass
    return total


def _evict_lora_cache_for_quota(needed: int) -> None:
    """Evict oldest unreferenced cache files until *needed* bytes fit, or raise.

    Only deletes files that are not currently referenced by in-memory LoRA stores.
    Never crashes the process solely because the cache is full — rejects the new
    artifact instead.
    """
    if needed <= 0:
        return
    usage = _lora_cache_usage_bytes()
    if usage + needed <= _H3_MAX_LORA_CACHE_BYTES:
        return
    # Build set of keys currently held in RAM (do not evict those).
    live = set()
    for k in list(_CUSTOM_LoRA.keys()):
        if isinstance(k, tuple) and len(k) >= 2:
            live.add(re.sub(r"[^A-Za-z0-9._-]", "_", f"{k[0].replace('/', '__')}__{k[1]}"))
    candidates = []
    try:
        for name in os.listdir(_LORA_CACHE_DIR):
            if name.endswith(".tmp") or name.startswith("."):
                continue
            p = os.path.join(_LORA_CACHE_DIR, name)
            if not os.path.isfile(p):
                continue
            # Skip if name is a prefix match of a live key
            if any(name.startswith(lk) or lk.startswith(name.split("__rev_")[0]) for lk in live):
                continue
            try:
                candidates.append((os.path.getmtime(p), os.path.getsize(p), p, name))
            except OSError:
                pass
    except OSError:
        pass
    candidates.sort()  # oldest first
    freed = 0
    for _mtime, size, p, name in candidates:
        if usage + needed - freed <= _H3_MAX_LORA_CACHE_BYTES:
            break
        try:
            os.remove(p)
            freed += size
            _log(f"LoRA cache evict (quota): {name} ({size / 1e6:.1f} MB)")
        except OSError as _e:
            _log(f"LoRA cache evict failed for {name}: {_e}", "WARN")
    if usage + needed - freed > _H3_MAX_LORA_CACHE_BYTES:
        raise gr.Error(
            f"LoRA cache quota exceeded (need {needed / 1e6:.0f} MB, "
            f"limit {_H3_MAX_LORA_CACHE_BYTES / 1e9:.1f} GB). "
            "Try again later or use a smaller LoRA."
        )