molgate_ui/data_cache.py

"""
Shared on-disk cache under data/ — network downloads land here for reuse.

Layout (see molgate_ui/contracts/STORAGE.md):
  data/pdb/{PDBID}.pdb     — PDB 快取（跨 session）
  data/sessions/{id}/      — 單次 run 工作區（非快取，但會引用快取）
"""

from __future__ import annotations

import shutil
from pathlib import Path

from molgate_ui.paths import PDB_DIR, SESSIONS, ensure_runtime_dirs


def cache_pdb(pdb_id: str, source: Path) -> Path:
    """Copy *source* into data/pdb/{PDBID}.pdb (create/update shared cache)."""
    ensure_runtime_dirs()
    pid = pdb_id.strip().upper()
    dest = PDB_DIR / f"{pid}.pdb"
    if not source.is_file():
        raise FileNotFoundError(source)
    dest.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(source, dest)
    return dest


def find_cached_pdb(pdb_id: str) -> Path | None:
    """Return data/pdb/{PDBID}.pdb if present."""
    pid = pdb_id.strip().upper()
    for name in (f"{pid}.pdb", f"{pid.lower()}.pdb"):
        p = PDB_DIR / name
        if p.is_file() and p.stat().st_size > 0:
            return p
    return None


def find_pdb_for_session(pdb_id: str, session_dir: Path) -> Path | None:
    """Resolve PDB: current session → shared cache → other sessions."""
    pid = pdb_id.strip().upper()
    session_path = session_dir / "pdb" / f"{pid}.pdb"
    if session_path.is_file() and session_path.stat().st_size > 0:
        return session_path

    cached = find_cached_pdb(pid)
    if cached:
        return cached

    if SESSIONS.is_dir():
        for p in sorted(
            SESSIONS.glob(f"*/pdb/{pid}.pdb"),
            key=lambda x: x.stat().st_mtime,
            reverse=True,
        ):
            if p.is_file() and p.stat().st_size > 0:
                return p
    return None


def write_session_pdb(pdb_id: str, session_dir: Path, text: str) -> Path:
    """Write PDB text into session and refresh shared cache."""
    pid = pdb_id.strip().upper()
    out = session_dir / "pdb" / f"{pid}.pdb"
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(text, encoding="utf-8")
    cache_pdb(pid, out)
    return out