Jyotisha/scripts/api_scratch_retention.py

"""Retention rules for scratch/local caches.

Counts only. Never prints file names or cache payloads.
"""
from __future__ import annotations

import json
import math
import os
import re
import secrets
import threading
import time
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Callable

ROOT = Path(__file__).resolve().parents[1]
UNREADABLE_KEEP_SECONDS = 24 * 60 * 60
_HEX64_JSON = re.compile(r"^[0-9a-f]{64}\.json$")
_ARTIFACT_JSON = re.compile(r"^[A-Za-z0-9_]+-[0-9a-f]{12}-[0-9a-f]{12}\.json$")
_GATEWAY_JOB_JSON = re.compile(r"^vgw_[0-9a-f]{20}\.json$")
_PRUNE_THREAD_NAME = "jyotish-scratch-prune"

# Directory names under scratch/local that this module already knows how to prune.
# Subdirectories of vedastro_adapter are listed so a future writer is not a new top-level dir.
REGISTERED_DIR_NAMES = frozenset(
    {
        "api_chart_cache",
        "high_rigor_jobs",
        "api_chart_response_jobs",
        "async_jobs.sqlite3",
        "vedastro_adapter",
        "vedastro_snapshot_cache",
        "vedastro_snapshot_negative_cache",
        "vedastro_gateway_jobs",
        "varga_segments",
    }
)

# Offline tools or caches that already prune themselves. Not a new retention rule.
OFFLINE_SCRATCH_NAMES = frozenset(
    {
        "consult_full_data_cache",
        # 与 consult_full_data_cache 同走 get_or_build，每次调用先 prune_cache(name=…) 删两天前的文件，自己会清；含出生资料，保留期 2 天。
        "report_reference_cache",
        "ocr_extract",
        "three_engine_parity",
        "public_release_deny_patterns.txt",
        "pdf_review_123456",
    }
)

REGISTRY: tuple[dict[str, Any], ...] = (
    {
        "name": "api_chart_cache",
        "holds_birth_data": True,
        "deletable": "created_at plus chart TTL, or unreadable json older than 24h",
    },
    {
        "name": "high_rigor_jobs",
        "holds_birth_data": True,
        "deletable": "expires_at_unix has passed; bad json is kept",
    },
    {
        "name": "api_chart_response_jobs",
        "holds_birth_data": True,
        "deletable": "expires_at_unix has passed; bad json is kept",
    },
    {
        "name": "async_jobs.sqlite3",
        "holds_birth_data": True,
        "deletable": "sqlite rows whose expires_at has passed",
    },
    {
        "name": "vedastro_adapter",
        "holds_birth_data": True,
        "deletable": "audit artifacts, response_cache, and official_full_snapshot_cache",
    },
    {
        "name": "vedastro_snapshot_cache",
        "holds_birth_data": True,
        "deletable": "reference_date before today UTC minus MAX_STALE_DAYS minus 1",
    },
    {
        "name": "vedastro_snapshot_negative_cache",
        "holds_birth_data": True,
        "deletable": "reference_date before today UTC",
    },
    {
        "name": "vedastro_gateway_jobs",
        "holds_birth_data": True,
        "deletable": "completed or failed jobs older than the keep window, plus their raw archive",
    },
    {
        "name": "varga_segments",
        "holds_birth_data": True,
        "deletable": "mtime older than JYOTISH_VARGA_SEGMENTS_KEEP_DAYS",
    },
)


def write_private_json(path: Path, payload: Any) -> None:
    """Atomic UTF-8 JSON, mode 0600. Compact, stable key order."""
    path.parent.mkdir(parents=True, exist_ok=True)
    temp = path.with_name(f"{path.name}.{secrets.token_hex(8)}.tmp")
    temp.write_text(
        json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":")),
        encoding="utf-8",
    )
    try:
        os.chmod(temp, 0o600)
    except OSError:
        pass
    os.replace(temp, path)
    try:
        os.chmod(path, 0o600)
    except OSError:
        pass


def _finite_float(value: Any) -> float | None:
    if isinstance(value, bool) or value is None:
        return None
    if isinstance(value, (int, float)):
        number = float(value)
    elif isinstance(value, str):
        text = value.strip()
        if not text:
            return None
        try:
            number = float(text)
        except ValueError:
            return None
    else:
        return None
    if not math.isfinite(number):
        return None
    return number


def _exact_int(value: Any) -> int | None:
    if isinstance(value, bool) or value is None:
        return None
    if isinstance(value, int):
        return value
    if isinstance(value, float):
        if not math.isfinite(value) or not value.is_integer():
            return None
        return int(value)
    if isinstance(value, str) and re.fullmatch(r"[+-]?\d+", value.strip()):
        return int(value.strip())
    return None


def normalize_chart_cache_birth(body: dict | None) -> dict[str, Any]:
    """Cache-key birth fields. Unparseable values stay as they arrived."""
    payload = body if isinstance(body, dict) else {}
    normalized: dict[str, Any] = {}
    for key in ("year", "month", "day"):
        parsed = _exact_int(payload.get(key))
        normalized[key] = payload.get(key) if parsed is None else parsed
    for key in ("hour", "minute", "lat", "lon", "tz"):
        raw = payload.get(key)
        parsed = _finite_float(raw)
        normalized[key] = raw if parsed is None else parsed
    second_raw = payload.get("second", 0)
    second = _finite_float(second_raw)
    normalized["second"] = second_raw if second is None else second
    return normalized


def normalize_chart_cache_day(value: Any) -> Any:
    if value is None or value == "":
        return value
    return str(value)[:10]


def _blank() -> dict[str, int]:
    return {"scanned": 0, "removed": 0, "errors": 0}


def _older_than(path: Path, moment: float, seconds: float) -> bool:
    try:
        return moment - path.stat().st_mtime > seconds
    except OSError:
        return False


def _unlink(path: Path, stats: dict[str, int]) -> None:
    try:
        path.unlink()
        stats["removed"] += 1
    except OSError:
        stats["errors"] += 1


def _repo_root() -> Path:
    return ROOT


def _api_module(api_module=None):
    if api_module is not None:
        return api_module
    import jyotish_api_server as api

    return api


def prune_async_job_records(now: float | None = None, api_module=None) -> dict[str, Any]:
    """Same file and sqlite rules as the old startup cleanup. Bad job JSON is kept."""
    api = _api_module(api_module)
    moment = time.time() if now is None else float(now)
    scanned = 0
    removed = 0
    errors = 0
    directories: dict[str, dict[str, int]] = {}
    if api._async_job_backend() == "sqlite":
        try:
            with api._sqlite_job_connection() as connection:
                scanned = int(connection.execute("SELECT COUNT(*) FROM async_jobs").fetchone()[0])
                deleted = connection.execute(
                    "DELETE FROM async_jobs WHERE expires_at IS NOT NULL AND expires_at <= ?",
                    (moment,),
                ).rowcount
                removed = int(deleted or 0)
        except OSError:
            errors += 1
        directories["high_rigor_jobs"] = _blank()
        directories["api_chart_response_jobs"] = _blank()
        directories["async_jobs_sqlite"] = {"scanned": scanned, "removed": removed, "errors": errors}
        return {
            "scope": "async_job_cleanup",
            "scanned": scanned,
            "removed": removed,
            "errors": errors,
            "directories": directories,
        }

    for scope, label in (
        (api._HIGH_RIGOR_JOB_SCOPE, "high_rigor_jobs"),
        (api._API_CHART_CACHE_SCOPE, "api_chart_response_jobs"),
    ):
        stats = _blank()
        try:
            directory = api._async_job_dir(scope)
        except OSError:
            stats["errors"] += 1
            directories[label] = stats
            errors += stats["errors"]
            continue
        if directory.is_dir():
            for path in directory.glob("*.json"):
                stats["scanned"] += 1
                try:
                    record = json.loads(path.read_text(encoding="utf-8"))
                except json.JSONDecodeError:
                    continue
                except OSError:
                    stats["errors"] += 1
                    continue
                expires_at = record.get("expires_at_unix") if isinstance(record, dict) else None
                if isinstance(expires_at, (int, float)) and not isinstance(expires_at, bool) and moment >= float(expires_at):
                    _unlink(path, stats)
        directories[label] = stats
        scanned += stats["scanned"]
        removed += stats["removed"]
        errors += stats["errors"]
    return {
        "scope": "async_job_cleanup",
        "scanned": scanned,
        "removed": removed,
        "errors": errors,
        "directories": directories,
    }


def _prune_chart_cache(moment: float, api_module=None) -> dict[str, int]:
    api = _api_module(api_module)
    stats = _blank()
    try:
        directory = api._api_chart_cache_dir()
        ttl = api._api_chart_cache_ttl_seconds()
    except OSError:
        stats["errors"] += 1
        return stats
    if not directory.is_dir():
        return stats
    for path in directory.glob("*.json"):
        if not _HEX64_JSON.match(path.name):
            continue
        stats["scanned"] += 1
        try:
            record = json.loads(path.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            if _older_than(path, moment, UNREADABLE_KEEP_SECONDS):
                _unlink(path, stats)
            continue
        except OSError:
            stats["errors"] += 1
            continue
        created_at = record.get("created_at") if isinstance(record, dict) else None
        if not isinstance(created_at, (int, float)) or isinstance(created_at, bool):
            if _older_than(path, moment, UNREADABLE_KEEP_SECONDS):
                _unlink(path, stats)
            continue
        if moment - float(created_at) > ttl:
            _unlink(path, stats)
    return stats


def _adapter():
    import vedastro_service_adapter as adapter

    return adapter


def _prune_created_at_dir(directory: Path, moment: float, ttl_for_record: Callable[[dict], float]) -> dict[str, int]:
    stats = _blank()
    if not directory.is_dir():
        return stats
    for path in directory.glob("*.json"):
        if not _HEX64_JSON.match(path.name):
            continue
        stats["scanned"] += 1
        try:
            record = json.loads(path.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            if _older_than(path, moment, UNREADABLE_KEEP_SECONDS):
                _unlink(path, stats)
            continue
        except OSError:
            stats["errors"] += 1
            continue
        created_at = record.get("created_at") if isinstance(record, dict) else None
        if not isinstance(created_at, (int, float)) or isinstance(created_at, bool):
            if _older_than(path, moment, UNREADABLE_KEEP_SECONDS):
                _unlink(path, stats)
            continue
        try:
            ttl = float(ttl_for_record(record))
        except (TypeError, ValueError):
            stats["errors"] += 1
            continue
        if moment - float(created_at) > ttl:
            _unlink(path, stats)
    return stats


def _prune_response_cache(moment: float) -> dict[str, int]:
    adapter = _adapter()
    return _prune_created_at_dir(
        Path(adapter.ARTIFACT_DIR) / "response_cache",
        moment,
        lambda _record: adapter._cache_ttl_seconds(),
    )


def _prune_snapshot_semantic_cache(moment: float) -> dict[str, int]:
    adapter = _adapter()

    def ttl_for(record: dict) -> float:
        if record.get("negative") is True:
            return adapter._official_full_snapshot_negative_ttl_seconds()
        return adapter._official_full_snapshot_cache_ttl_seconds()

    return _prune_created_at_dir(
        Path(adapter.ARTIFACT_DIR) / "official_full_snapshot_cache",
        moment,
        ttl_for,
    )


def _audit_keep_seconds() -> float:
    raw = os.environ.get("VEDASTRO_AUDIT_ARTIFACT_KEEP_DAYS", "7").strip()
    try:
        days = float(raw) if raw else 7.0
    except ValueError:
        days = 7.0
    return max(0.0, days) * 86400.0


def _prune_audit_artifacts(moment: float) -> dict[str, int]:
    adapter = _adapter()
    stats = _blank()
    directory = Path(adapter.ARTIFACT_DIR)
    if not directory.is_dir():
        return stats
    enabled = adapter._audit_artifacts_enabled()
    max_age = _audit_keep_seconds() if enabled else UNREADABLE_KEEP_SECONDS
    for path in directory.glob("*.json"):
        if not path.is_file() or not _ARTIFACT_JSON.match(path.name):
            continue
        stats["scanned"] += 1
        if _older_than(path, moment, max_age):
            _unlink(path, stats)
    return stats


def _parse_day(value: Any):
    if not isinstance(value, str):
        return None
    raw = value[:10]
    try:
        return datetime.strptime(raw, "%Y-%m-%d").date()
    except ValueError:
        return None


def _prune_reference_dated(directory: Path, moment: float, *, earliest_kept) -> dict[str, int]:
    stats = _blank()
    if not directory.is_dir():
        return stats
    for path in directory.glob("*.json"):
        if not _HEX64_JSON.match(path.name):
            continue
        stats["scanned"] += 1
        try:
            record = json.loads(path.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            if _older_than(path, moment, UNREADABLE_KEEP_SECONDS):
                _unlink(path, stats)
            continue
        except OSError:
            stats["errors"] += 1
            continue
        ref = _parse_day(record.get("reference_date") if isinstance(record, dict) else None)
        if ref is None:
            if _older_than(path, moment, UNREADABLE_KEEP_SECONDS):
                _unlink(path, stats)
            continue
        if ref < earliest_kept:
            _unlink(path, stats)
    return stats


def _prune_vedastro_snapshot_cache(moment: float) -> dict[str, int]:
    from vedastro_snapshot_cache import MAX_STALE_DAYS

    today = datetime.fromtimestamp(moment, timezone.utc).date()
    earliest = today - timedelta(days=int(MAX_STALE_DAYS) + 1)
    raw = os.environ.get("JYOTISH_VEDASTRO_SNAPSHOT_CACHE_DIR", "").strip()
    directory = Path(raw) if raw else _repo_root() / "scratch" / "local" / "vedastro_snapshot_cache"
    return _prune_reference_dated(directory, moment, earliest_kept=earliest)


def _prune_vedastro_negative_cache(moment: float) -> dict[str, int]:
    today = datetime.fromtimestamp(moment, timezone.utc).date()
    raw = os.environ.get("JYOTISH_VEDASTRO_SNAPSHOT_NEGATIVE_CACHE_DIR", "").strip()
    directory = Path(raw) if raw else _repo_root() / "scratch" / "local" / "vedastro_snapshot_negative_cache"
    return _prune_reference_dated(directory, moment, earliest_kept=today)


def _parse_updated_at(value: Any) -> float | None:
    if not isinstance(value, str) or not value.strip():
        return None
    text = value.strip().replace("Z", "+00:00")
    try:
        parsed = datetime.fromisoformat(text)
    except ValueError:
        return None
    if parsed.tzinfo is None:
        parsed = parsed.replace(tzinfo=timezone.utc)
    return parsed.timestamp()


def _gateway_keep_seconds() -> float:
    raw = os.environ.get("VEDASTRO_GATEWAY_JOB_KEEP_SECONDS", "86400").strip()
    try:
        return max(0.0, float(raw) if raw else 86400.0)
    except ValueError:
        return 86400.0


def _prune_gateway_jobs(moment: float) -> dict[str, int]:
    stats = _blank()
    raw = os.environ.get("VEDASTRO_GATEWAY_QUEUE_DIR", "").strip()
    directory = Path(raw).expanduser() if raw else _repo_root() / "scratch" / "local" / "vedastro_gateway_jobs"
    if not directory.is_dir():
        return stats
    keep = _gateway_keep_seconds()
    for path in directory.glob("*.json"):
        if not _GATEWAY_JOB_JSON.match(path.name):
            continue
        stats["scanned"] += 1
        try:
            record = json.loads(path.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            continue
        except OSError:
            stats["errors"] += 1
            continue
        if not isinstance(record, dict):
            continue
        status = str(record.get("status") or "")
        if status not in {"completed", "failed"}:
            continue
        updated = _parse_updated_at(record.get("updated_at"))
        if updated is None or moment - updated <= keep:
            continue
        companion = directory / f"{path.stem}.official_raw_response.json"
        _unlink(path, stats)
        if companion.is_file():
            stats["scanned"] += 1
            _unlink(companion, stats)
    return stats


def _varga_keep_seconds() -> float:
    raw = os.environ.get("JYOTISH_VARGA_SEGMENTS_KEEP_DAYS", "7").strip()
    try:
        days = float(raw) if raw else 7.0
    except ValueError:
        days = 7.0
    return max(0.0, days) * 86400.0


def _prune_varga_segments(moment: float) -> dict[str, int]:
    stats = _blank()
    raw = os.environ.get("JYOTISH_CACHE_DIR", "").strip()
    root = Path(raw) if raw else _repo_root() / "scratch" / "local"
    directory = root / "varga_segments"
    if not directory.is_dir():
        return stats
    max_age = _varga_keep_seconds()
    for path in directory.glob("*.json"):
        if not _HEX64_JSON.match(path.name):
            continue
        stats["scanned"] += 1
        if _older_than(path, moment, max_age):
            _unlink(path, stats)
    return stats


def _stamp(stats: dict[str, int], started: float) -> dict[str, Any]:
    return {
        "scanned": int(stats.get("scanned", 0) or 0),
        "removed": int(stats.get("removed", 0) or 0),
        "errors": int(stats.get("errors", 0) or 0),
        "elapsed_ms": round((time.perf_counter() - started) * 1000, 3),
    }


def prune_all(now: float | None = None, api_module=None) -> dict[str, Any]:
    """Prune every registered directory. Exceptions become error counts."""
    moment = time.time() if now is None else float(now)
    api = _api_module(api_module)
    started = time.perf_counter()
    directories: dict[str, dict[str, Any]] = {}

    def _run(name: str, func: Callable[[float], dict[str, int]]) -> None:
        t0 = time.perf_counter()
        try:
            stats = func(moment)
        except Exception:
            stats = {"scanned": 0, "removed": 0, "errors": 1}
        directories[name] = _stamp(stats, t0)

    _run("api_chart_cache", lambda moment: _prune_chart_cache(moment, api))
    job_started = time.perf_counter()
    try:
        job = prune_async_job_records(moment, api)
        job_dirs = job.get("directories") or {}
    except Exception:
        job_dirs = {
            "high_rigor_jobs": {"scanned": 0, "removed": 0, "errors": 1},
            "api_chart_response_jobs": _blank(),
        }
    job_elapsed = round((time.perf_counter() - job_started) * 1000, 3)
    for name, stats in job_dirs.items():
        stamped = _stamp(stats, job_started)
        stamped["elapsed_ms"] = job_elapsed
        directories[name] = stamped
    _run("vedastro_response_cache", _prune_response_cache)
    _run("vedastro_official_full_snapshot_cache", _prune_snapshot_semantic_cache)
    _run("vedastro_audit_artifacts", _prune_audit_artifacts)
    _run("vedastro_snapshot_cache", _prune_vedastro_snapshot_cache)
    _run("vedastro_snapshot_negative_cache", _prune_vedastro_negative_cache)
    _run("vedastro_gateway_jobs", _prune_gateway_jobs)
    _run("varga_segments", _prune_varga_segments)
    return {
        "scanned": sum(item["scanned"] for item in directories.values()),
        "removed": sum(item["removed"] for item in directories.values()),
        "errors": sum(item["errors"] for item in directories.values()),
        "elapsed_ms": round((time.perf_counter() - started) * 1000, 3),
        "directories": directories,
    }


def start_periodic_prune(interval_s: float | None = None):
    """Daemon thread. interval <= 0 starts nothing. Startup prune is separate."""
    if interval_s is None:
        raw = os.environ.get("JYOTISH_SCRATCH_PRUNE_INTERVAL_SECONDS", "1800").strip()
        try:
            interval_s = float(raw) if raw else 1800.0
        except ValueError:
            interval_s = 1800.0
    if interval_s <= 0:
        return None

    def _loop() -> None:
        while True:
            time.sleep(interval_s)
            try:
                prune_all()
            except Exception:
                continue

    thread = threading.Thread(target=_loop, name=_PRUNE_THREAD_NAME, daemon=True)
    thread.start()
    return thread