fithia2/libs/backtest/snapshots.py

"""Snapshot registry, alias resolution, and canonical path helpers."""
from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from libs.common.config import get_settings

_REGISTRY_PATH = Path("configs/snapshots/registry.json")
_LEGACY_SNAPSHOT_FALLBACK_ROOT = Path("data/datasets/snapshots")


class UnknownSnapshotError(ValueError):
    """Raised when a snapshot id cannot be resolved."""


@dataclass(frozen=True)
class SnapshotResolution:
    requested_snapshot_id: str
    canonical_snapshot_id: str
    refresh_policy: str | None
    purpose: str | None
    expected_feature_set: str | None
    is_alias: bool
    is_registry_managed: bool
    is_legacy_passthrough: bool


def _load_registry_payload() -> dict[str, Any]:
    if not _REGISTRY_PATH.exists():
        raise FileNotFoundError(f"Snapshot registry not found: {_REGISTRY_PATH}")
    return json.loads(_REGISTRY_PATH.read_text())


def load_snapshot_registry() -> dict[str, Any]:
    """Return the raw snapshot registry payload."""
    return _load_registry_payload()


def get_canonical_spec(canonical_snapshot_id: str) -> dict[str, Any]:
    """Return the registry spec for a managed canonical snapshot."""
    payload = _load_registry_payload()
    spec = payload.get("canonicals", {}).get(canonical_snapshot_id)
    if spec is None:
        raise UnknownSnapshotError(f"Unknown canonical snapshot id: {canonical_snapshot_id}")
    return dict(spec)


def get_requested_aliases(canonical_snapshot_id: str) -> list[str]:
    """Return all legacy aliases that resolve to *canonical_snapshot_id*."""
    payload = _load_registry_payload()
    aliases = payload.get("aliases", {})
    return sorted(
        alias
        for alias, alias_spec in aliases.items()
        if alias_spec.get("canonical_id") == canonical_snapshot_id
    )


def _snapshot_root_candidates(snapshot_dir: str | Path | None = None) -> list[Path]:
    if snapshot_dir is not None:
        return [Path(snapshot_dir)]
    settings = get_settings()
    roots = [Path(settings.parquet_dir), _LEGACY_SNAPSHOT_FALLBACK_ROOT]
    deduped: list[Path] = []
    seen: set[Path] = set()
    for root in roots:
        root = root.resolve()
        if root in seen:
            continue
        seen.add(root)
        deduped.append(root)
    return deduped


def snapshot_exists(snapshot_id: str, snapshot_dir: str | Path | None = None) -> bool:
    """Return True when the raw snapshot directory exists in any known root."""
    for root in _snapshot_root_candidates(snapshot_dir):
        if (root / snapshot_id).exists():
            return True
    return False


def resolve_snapshot(
    requested_snapshot_id: str,
    *,
    snapshot_dir: str | Path | None = None,
    allow_legacy_passthrough: bool = True,
) -> SnapshotResolution:
    """Resolve a requested snapshot id to its canonical snapshot id.

    Registry-managed aliases map to one of the two canonicals. For legacy
    directories outside the managed alias set, phase-1 keeps a passthrough
    fallback so old experiments do not break immediately.
    """
    payload = _load_registry_payload()
    canonicals = payload.get("canonicals", {})
    aliases = payload.get("aliases", {})

    if requested_snapshot_id in canonicals:
        spec = canonicals[requested_snapshot_id]
        return SnapshotResolution(
            requested_snapshot_id=requested_snapshot_id,
            canonical_snapshot_id=requested_snapshot_id,
            refresh_policy=spec.get("refresh_policy"),
            purpose=spec.get("purpose"),
            expected_feature_set=None,
            is_alias=False,
            is_registry_managed=True,
            is_legacy_passthrough=False,
        )

    if requested_snapshot_id in aliases:
        alias_spec = aliases[requested_snapshot_id]
        canonical_snapshot_id = str(alias_spec["canonical_id"])
        canonical_spec = canonicals.get(canonical_snapshot_id, {})
        return SnapshotResolution(
            requested_snapshot_id=requested_snapshot_id,
            canonical_snapshot_id=canonical_snapshot_id,
            refresh_policy=canonical_spec.get("refresh_policy"),
            purpose=alias_spec.get("purpose") or canonical_spec.get("purpose"),
            expected_feature_set=alias_spec.get("expected_feature_set"),
            is_alias=True,
            is_registry_managed=True,
            is_legacy_passthrough=False,
        )

    if allow_legacy_passthrough and snapshot_exists(requested_snapshot_id, snapshot_dir=snapshot_dir):
        return SnapshotResolution(
            requested_snapshot_id=requested_snapshot_id,
            canonical_snapshot_id=requested_snapshot_id,
            refresh_policy="legacy_passthrough",
            purpose="legacy_passthrough",
            expected_feature_set=None,
            is_alias=False,
            is_registry_managed=False,
            is_legacy_passthrough=True,
        )

    raise UnknownSnapshotError(f"Unknown snapshot id: {requested_snapshot_id}")


def resolve_snapshot_path(
    requested_snapshot_id: str,
    *,
    snapshot_dir: str | Path | None = None,
    allow_legacy_passthrough: bool = True,
) -> Path | None:
    """Resolve the on-disk snapshot directory.

    Prefer the canonical directory. If it does not exist yet, fall back to the
    requested legacy alias directory so phase-1 remains parity-safe until the
    canonical snapshot has been rebuilt.
    """
    resolution = resolve_snapshot(
        requested_snapshot_id,
        snapshot_dir=snapshot_dir,
        allow_legacy_passthrough=allow_legacy_passthrough,
    )
    candidate_ids = [resolution.canonical_snapshot_id]
    if resolution.requested_snapshot_id != resolution.canonical_snapshot_id:
        candidate_ids.append(resolution.requested_snapshot_id)

    seen: set[Path] = set()
    for root in _snapshot_root_candidates(snapshot_dir):
        for snapshot_id in candidate_ids:
            candidate = (root / snapshot_id).resolve()
            if candidate in seen:
                continue
            seen.add(candidate)
            if candidate.exists():
                return candidate
    return None


def canonical_snapshot_dir(canonical_snapshot_id: str) -> Path:
    """Return the managed canonical directory under data/parquet."""
    settings = get_settings()
    return Path(settings.parquet_dir).resolve() / canonical_snapshot_id