md_to_gost/md2gost/diagram_includes.py

"""HTTP(S) include cache: index URL → local file, never overwrite existing cache files."""

from __future__ import annotations

import hashlib
import json
import logging
import re
import sys
from pathlib import Path
from urllib.parse import urlparse, urljoin

import requests

_log = logging.getLogger(__name__)

INCLUDE_CACHE_INDEX = "md2gost.include-cache.json"
INCLUDE_CACHE_DIR = "include-cache"
MAX_INCLUDE_DEPTH = 8

_HTTP_RE = re.compile(r"^https?://", re.I)
_INCLUDE_LINE_RE = re.compile(
    r"^[ \t]*!(?:includeurl|include)[ \t]+(?P<q>[\"']?)(?P<ref>[^\s\"']+)(?P=q)",
    re.I | re.M,
)


def app_dir() -> Path:
    """Directory next to the running application (exe or cwd for python -m)."""
    if getattr(sys, "frozen", False):
        return Path(sys.executable).resolve().parent
    return Path.cwd()


def include_cache_index_path(base: Path | None = None) -> Path:
    return (base or app_dir()) / INCLUDE_CACHE_INDEX


def include_cache_files_dir(base: Path | None = None) -> Path:
    return (base or app_dir()) / INCLUDE_CACHE_DIR


def is_http_url(value: str) -> bool:
    return bool(_HTTP_RE.match((value or "").strip()))


def _load_index(path: Path) -> dict[str, str]:
    if not path.is_file():
        return {}
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        _log.warning("Не удалось прочитать индекс кэша includes: %s", exc)
        return {}
    if not isinstance(data, dict):
        return {}
    out: dict[str, str] = {}
    for key, val in data.items():
        if isinstance(key, str) and isinstance(val, str) and is_http_url(key):
            out[key] = val
    return out


def _save_index(path: Path, index: dict[str, str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(index, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )


def _unique_cache_path(url: str, files_dir: Path) -> Path:
    """New file path for URL; never reuse/overwrite an existing cache file."""
    files_dir.mkdir(parents=True, exist_ok=True)
    parsed = urlparse(url)
    base = Path(parsed.path).name or "include.puml"
    if not base.lower().endswith((".puml", ".iuml", ".pu", ".txt")):
        base = base + ".puml" if "." not in base else base
    digest = hashlib.sha256(url.encode("utf-8")).hexdigest()[:10]
    stem = Path(base).stem
    suffix = Path(base).suffix or ".puml"
    # sanitize stem
    safe = re.sub(r"[^a-zA-Z0-9._-]+", "_", stem)[:40] or "include"
    candidate = files_dir / f"{safe}_{digest}{suffix}"
    if not candidate.exists():
        return candidate
    n = 1
    while True:
        alt = files_dir / f"{safe}_{digest}_{n}{suffix}"
        if not alt.exists():
            return alt
        n += 1


def _download(url: str, dest: Path) -> None:
    dest.parent.mkdir(parents=True, exist_ok=True)
    tmp = dest.with_suffix(dest.suffix + ".part")
    try:
        resp = requests.get(url, timeout=60)
        resp.raise_for_status()
        tmp.write_bytes(resp.content)
        tmp.replace(dest)
    except Exception:
        tmp.unlink(missing_ok=True)
        raise


def resolve_url(
    url: str,
    *,
    base: Path | None = None,
    depth: int = 0,
) -> Path:
    """
    Return local path for URL. Uses index if present and file exists;
    otherwise downloads to a *new* file and records the mapping.
    Never overwrites an existing cache file for a different URL.
    """
    url = url.strip()
    if not is_http_url(url):
        raise ValueError(f"Не URL: {url}")
    if depth > MAX_INCLUDE_DEPTH:
        raise RuntimeError(f"Слишком глубокая цепочка includes ({MAX_INCLUDE_DEPTH}): {url}")

    root = base or app_dir()
    index_path = include_cache_index_path(root)
    files_dir = include_cache_files_dir(root)
    index = _load_index(index_path)

    cached = index.get(url)
    if cached:
        path = Path(cached)
        if path.is_file():
            _rewrite_nested_includes(path, parent_url=url, base=root, depth=depth)
            return path
        _log.warning("Файл кэша пропал (%s), скачаю заново: %s", path, url)

    dest = _unique_cache_path(url, files_dir)
    try:
        _download(url, dest)
    except Exception as exc:
        raise RuntimeError(
            f"Не удалось скачать include {url}: {exc}. "
            "Нет сети и нет кэша — положите файл вручную или проверьте URL."
        ) from exc

    index[url] = str(dest.resolve())
    _save_index(index_path, index)
    _log.info("Include cached: %s → %s", url, dest)
    _rewrite_nested_includes(dest, parent_url=url, base=root, depth=depth)
    return dest


def _absolute_include_url(ref: str, parent_url: str | None) -> str | None:
    """Turn include ref into absolute http(s) URL, or None if local/stdlib."""
    ref = ref.strip()
    if is_http_url(ref):
        return ref
    if ref.startswith("<") and ref.endswith(">"):
        return None  # PlantUML stdlib
    if parent_url and not Path(ref).is_absolute():
        # relative to parent URL directory
        return urljoin(parent_url, ref)
    return None


def _rewrite_nested_includes(
    path: Path, *, parent_url: str, base: Path, depth: int,
) -> None:
    """Rewrite http(s) / relative !include inside a cached .puml to local paths."""
    try:
        text = path.read_text(encoding="utf-8")
    except OSError:
        return

    changed = False

    def repl(match: re.Match[str]) -> str:
        nonlocal changed
        ref = match.group("ref")
        abs_url = _absolute_include_url(ref, parent_url)
        if not abs_url:
            return match.group(0)
        local = resolve_url(abs_url, base=base, depth=depth + 1)
        changed = True
        local_s = str(local.resolve()).replace("\\", "/")
        return f"!include {local_s}"

    new_text = _INCLUDE_LINE_RE.sub(repl, text)
    if changed and new_text != text:
        path.write_text(new_text, encoding="utf-8")


def resolve_include_ref(
    ref: str,
    *,
    base: Path | None = None,
    schemes_dir: Path | None = None,
) -> str:
    """
    Resolve include entry to a local path string suitable for !include.
    Local paths are resolved relative to schemes_dir / diagrams / cwd.
    """
    ref = (ref or "").strip()
    if not ref:
        raise ValueError("Пустой include")
    if is_http_url(ref):
        return str(resolve_url(ref, base=base).resolve()).replace("\\", "/")

    path = Path(ref)
    if path.is_file():
        return str(path.resolve()).replace("\\", "/")
    candidates: list[Path] = []
    if schemes_dir:
        candidates.append(schemes_dir / ref)
    from . import package_dir
    candidates.append(Path(package_dir()) / "diagrams" / ref)
    candidates.append(Path.cwd() / ref)
    for cand in candidates:
        if cand.is_file():
            return str(cand.resolve()).replace("\\", "/")
    raise FileNotFoundError(f"Include не найден: {ref}")


def rewrite_http_includes_in_source(source: str, *, base: Path | None = None) -> str:
    """Replace http(s) !include / !includeurl in prepared source with cached local paths."""

    def repl(match: re.Match[str]) -> str:
        ref = match.group("ref")
        if not is_http_url(ref):
            return match.group(0)
        local = str(resolve_url(ref, base=base).resolve()).replace("\\", "/")
        return f"!include {local}"

    return _INCLUDE_LINE_RE.sub(repl, source)


def reset_include_cache(base: Path | None = None) -> int:
    """
    Delete index and only files listed in it. Returns number of files removed.
    Other files in include-cache/ are left untouched.
    """
    root = base or app_dir()
    index_path = include_cache_index_path(root)
    index = _load_index(index_path)
    removed = 0
    for _url, path_s in list(index.items()):
        path = Path(path_s)
        try:
            if path.is_file():
                path.unlink()
                removed += 1
        except OSError as exc:
            _log.warning("Не удалось удалить %s: %s", path, exc)
    if index_path.is_file():
        index_path.unlink(missing_ok=True)
    return removed


def cache_entry_count(base: Path | None = None) -> int:
    return len(_load_index(include_cache_index_path(base or app_dir())))