fithia2/libs/export/canonical_snapshots.py

"""Canonical snapshot build orchestration."""
from __future__ import annotations

import datetime as dt
import json
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any

import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.parquet as pq

from libs.backtest.snapshots import (
    canonical_snapshot_dir,
    get_canonical_spec,
    get_requested_aliases,
    resolve_snapshot,
)
from libs.backtest.snapshot_store import SnapshotStore
from libs.common.config import get_settings
from libs.common.logging import get_logger
from libs.common.time_utils import utc_now
from libs.db.session import get_session
from libs.export.snapshot_export import export_dataset_snapshot

logger = get_logger(__name__)


def _coerce_schema(table: pa.Table, target_schema: pa.Schema) -> pa.Table:
    """Cast columns in table to match target_schema types where they differ numerically."""
    for i in range(len(target_schema)):
        field = target_schema.field(i)
        if field.name not in table.schema.names:
            continue
        col_idx = table.schema.get_field_index(field.name)
        existing_type = table.schema.field(col_idx).type
        if existing_type == field.type:
            continue
        try:
            table = table.set_column(col_idx, field, table.column(field.name).cast(field.type, safe=False))
        except Exception:
            pass  # leave as-is; concat_tables will promote or raise with clear error
    return table

_ENRICHMENT_SCRIPT_BY_STEP = {
    "earnings_history_enrich": Path("scripts/enrich_earnings_history_features.py"),
    "peer_surprise_enrich": Path("scripts/enrich_peer_surprise_features.py"),
    "catalyst_persistence_enrich": Path("scripts/enrich_catalyst_persistence_features.py"),
    "tier2_enrich": Path("scripts/enrich_tier2_features.py"),
    "tier3_enrich": Path("scripts/enrich_tier3_features.py"),
    "technical_enrich": Path("scripts/enrich_technical_features.py"),
    "macro_enrich": Path("scripts/enrich_macro_features.py"),
    "prior_drift_enrich": Path("scripts/enrich_prior_drift.py"),
}


def _coverage_end_date(snapshot_dir: Path) -> str | None:
    max_date: dt.date | None = None
    for split_name in ("train", "valid", "test"):
        parquet_path = snapshot_dir / f"{split_name}.parquet"
        if not parquet_path.exists():
            continue
        table = pq.read_table(str(parquet_path))
        if "event_date" not in table.column_names:
            continue
        table = table.select(["event_date"])
        for value in table.column("event_date").to_pylist():
            if isinstance(value, str):
                candidate = dt.date.fromisoformat(value[:10])
            elif isinstance(value, dt.datetime):
                candidate = value.date()
            elif isinstance(value, dt.date):
                candidate = value
            else:
                continue
            if max_date is None or candidate > max_date:
                max_date = candidate
    return max_date.isoformat() if max_date is not None else None


def _write_canonical_manifest(
    snapshot_dir: Path,
    canonical_snapshot_id: str,
    spec: dict[str, Any],
    extra_fields: dict[str, Any] | None = None,
) -> None:
    manifest_path = snapshot_dir / "manifest.json"
    manifest = json.loads(manifest_path.read_text()) if manifest_path.exists() else {}
    manifest.update(
        {
            "snapshot_id": canonical_snapshot_id,
            "canonical_snapshot_id": canonical_snapshot_id,
            "output_dir": str(snapshot_dir.resolve()),
            "feature_sets": list(spec.get("feature_sets") or []),
            "enrichment_steps": list(spec.get("enrichment_steps") or []),
            "requested_aliases": get_requested_aliases(canonical_snapshot_id),
            "coverage_end_date": _coverage_end_date(snapshot_dir),
            "last_refresh_utc": utc_now().isoformat(),
            "refresh_policy": spec.get("refresh_policy"),
            "purpose": spec.get("purpose"),
            "expected_feature_columns": list(spec.get("expected_feature_columns") or []),
            "universe_profile": spec.get("universe_profile"),
            "start_date": spec.get("start_date"),
            "end_date": spec.get("end_date"),
        }
    )
    if extra_fields:
        manifest.update(extra_fields)
    manifest_path.write_text(json.dumps(manifest, indent=2))


def _run_enrichment_step(step_name: str, input_dir: Path, output_dir: Path) -> None:
    script_path = _ENRICHMENT_SCRIPT_BY_STEP[step_name]
    cmd = [
        sys.executable,
        str(script_path),
        "--input",
        str(input_dir),
        "--output",
        str(output_dir),
    ]
    logger.info("canonical_snapshot_enrichment_start", step=step_name, cmd=cmd)
    env = {**__import__("os").environ, "PYTHONUNBUFFERED": "1"}
    subprocess.run(cmd, check=True, env=env)


async def _materialize_runtime_backfills(snapshot_dir: Path) -> list[str]:
    settings = get_settings()
    return await SnapshotStore._async_materialize_snapshot_dir(
        snapshot_dir,
        oracle_url=settings.stock_oracle_url,
        db_dsn=settings.postgres_dsn,
    )


def _atomic_swap_directory(staged_final_dir: Path, target_dir: Path) -> None:
    backup_dir = target_dir.with_name(f"{target_dir.name}.bak.{utc_now().strftime('%Y%m%d%H%M%S%f')}")
    backup_created = False
    try:
        if target_dir.exists():
            target_dir.rename(backup_dir)
            backup_created = True
        staged_final_dir.rename(target_dir)
    except Exception:
        if backup_created and not target_dir.exists() and backup_dir.exists():
            backup_dir.rename(target_dir)
        raise
    else:
        if backup_created and backup_dir.exists():
            shutil.rmtree(backup_dir)


def _resolve_bootstrap_source_dir(source_snapshot_id: str) -> Path:
    """Resolve an existing snapshot dir for bootstrap from managed roots."""
    source_path = Path(source_snapshot_id)
    if source_path.is_absolute() and source_path.exists():
        return source_path

    settings = get_settings()
    candidate_roots = [
        Path(settings.parquet_dir),
        Path("data/datasets/snapshots"),
    ]
    for root in candidate_roots:
        candidate = root / source_snapshot_id
        if candidate.exists():
            return candidate.resolve()
    raise FileNotFoundError(f"Bootstrap source snapshot not found: {source_snapshot_id}")


async def build_canonical_snapshot(
    requested_snapshot_id: str,
    *,
    manual: bool = False,
) -> Path:
    """Build a registry-managed canonical snapshot via full rebuild + atomic swap."""
    resolution = resolve_snapshot(requested_snapshot_id, allow_legacy_passthrough=False)
    canonical_snapshot_id = resolution.canonical_snapshot_id
    spec = get_canonical_spec(canonical_snapshot_id)

    refresh_policy = str(spec.get("refresh_policy") or "")
    if refresh_policy == "manual_only" and not manual:
        raise RuntimeError(f"Snapshot '{canonical_snapshot_id}' is frozen for auto-refresh.")

    target_dir = canonical_snapshot_dir(canonical_snapshot_id)
    target_dir.parent.mkdir(parents=True, exist_ok=True)

    with tempfile.TemporaryDirectory(
        prefix=f".{canonical_snapshot_id}.stage.",
        dir=str(target_dir.parent),
    ) as stage_root_str:
        stage_root = Path(stage_root_str)
        current_dir = stage_root / "00_base"

        start_date = spec.get("start_date")
        end_date = spec.get("end_date")
        async with get_session() as session:
            await export_dataset_snapshot(
                session=session,
                snapshot_id=current_dir.name,
                split_policy="temporal_70_15_15",
                output_dir=stage_root,
                feature_versions=list(spec.get("feature_versions") or []),
                universe_profile=spec.get("universe_profile"),
                label_version=spec.get("label_version", "label-1.0.0"),
                start_date=dt.date.fromisoformat(start_date) if start_date else None,
                end_date=dt.date.fromisoformat(end_date) if end_date else None,
                include_export_enrichments=False,
            )

        steps = [step for step in list(spec.get("enrichment_steps") or []) if step != "base_export"]
        for idx, step in enumerate(steps, start=1):
            is_last = idx == len(steps)
            next_dir = stage_root / (canonical_snapshot_id if is_last else f"{idx:02d}_{step}")
            _run_enrichment_step(step, current_dir, next_dir)
            current_dir = next_dir

        if current_dir.name != canonical_snapshot_id:
            final_dir = stage_root / canonical_snapshot_id
            shutil.copytree(current_dir, final_dir)
            current_dir = final_dir

        materialized_columns = await _materialize_runtime_backfills(current_dir)
        _write_canonical_manifest(
            current_dir,
            canonical_snapshot_id,
            spec,
            extra_fields={
                "materialized_feature_columns": materialized_columns,
                "materialization_mode": "runtime_backfill_persisted_v1",
            },
        )
        _atomic_swap_directory(current_dir, target_dir)

    return target_dir


async def bootstrap_canonical_snapshot(
    requested_snapshot_id: str,
    *,
    source_snapshot_id: str,
    enrichment_steps: list[str] | None = None,
) -> Path:
    """Materialize a canonical snapshot from an existing local snapshot directory.

    This is a parity-preserving migration path for phase 1 when the legacy
    enriched snapshot is the current source of truth.
    """
    resolution = resolve_snapshot(requested_snapshot_id, allow_legacy_passthrough=False)
    canonical_snapshot_id = resolution.canonical_snapshot_id
    spec = get_canonical_spec(canonical_snapshot_id)

    source_dir = _resolve_bootstrap_source_dir(source_snapshot_id)

    target_dir = canonical_snapshot_dir(canonical_snapshot_id)
    target_dir.parent.mkdir(parents=True, exist_ok=True)

    with tempfile.TemporaryDirectory(
        prefix=f".{canonical_snapshot_id}.bootstrap.",
        dir=str(target_dir.parent),
    ) as stage_root_str:
        stage_root = Path(stage_root_str)
        current_dir = stage_root / canonical_snapshot_id
        shutil.copytree(source_dir, current_dir)

        for idx, step in enumerate(list(enrichment_steps or []), start=1):
            is_last = idx == len(list(enrichment_steps or []))
            next_dir = stage_root / (canonical_snapshot_id if is_last else f"bootstrap_{idx:02d}_{step}")
            _run_enrichment_step(step, current_dir, next_dir)
            current_dir = next_dir

        materialized_columns = await _materialize_runtime_backfills(current_dir)
        _write_canonical_manifest(
            current_dir,
            canonical_snapshot_id,
            spec,
            extra_fields={
                "bootstrap_source_snapshot_id": source_snapshot_id,
                "bootstrap_mode": "copy_existing_snapshot",
                "materialized_feature_columns": materialized_columns,
                "materialization_mode": "runtime_backfill_persisted_v1",
            },
        )
        _atomic_swap_directory(current_dir, target_dir)

    return target_dir


async def incremental_update_canonical_snapshot(
    requested_snapshot_id: str,
) -> Path:
    """Append only new events (since coverage_end_date) from DB to existing canonical.

    Does not rebuild existing rows. Runs enrichment scripts only on new rows.
    New rows are appended to the test split (they are always the most recent events).
    """
    resolution = resolve_snapshot(requested_snapshot_id, allow_legacy_passthrough=False)
    canonical_snapshot_id = resolution.canonical_snapshot_id
    spec = get_canonical_spec(canonical_snapshot_id)
    target_dir = canonical_snapshot_dir(canonical_snapshot_id)

    if not target_dir.exists():
        raise FileNotFoundError(f"Canonical snapshot not found: {target_dir}")

    # Determine start_date for new events
    max_date_str = _coverage_end_date(target_dir)
    if max_date_str is None:
        raise ValueError("Cannot determine coverage_end_date from existing canonical")
    new_start = dt.date.fromisoformat(max_date_str) + dt.timedelta(days=1)

    logger.info(
        "incremental_update_canonical_start",
        canonical_snapshot_id=canonical_snapshot_id,
        existing_coverage_end=max_date_str,
        new_start=str(new_start),
    )

    with tempfile.TemporaryDirectory(
        prefix=f".{canonical_snapshot_id}.incremental.",
        dir=str(target_dir.parent),
    ) as stage_root_str:
        stage_root = Path(stage_root_str)

        # 1. Export new events from DB (start_date = day after current coverage)
        end_date_str = spec.get("end_date")
        async with get_session() as session:
            await export_dataset_snapshot(
                session=session,
                snapshot_id="new_events",
                split_policy="temporal_70_15_15",
                output_dir=stage_root,
                feature_versions=list(spec.get("feature_versions") or []),
                universe_profile=spec.get("universe_profile"),
                label_version=spec.get("label_version", "label-1.0.0"),
                start_date=new_start,
                end_date=dt.date.fromisoformat(end_date_str) if end_date_str else None,
                include_export_enrichments=False,
            )

        new_export_dir = stage_root / "new_events"

        # Combine all exported splits into one batch (all new rows are recent → go to test)
        all_new_tables = []
        for split in ("train", "valid", "test"):
            p = new_export_dir / f"{split}.parquet"
            if p.exists():
                t = pq.read_table(str(p))
                if len(t) > 0:
                    all_new_tables.append(t)

        if not all_new_tables:
            logger.info("incremental_update_no_new_events", new_start=str(new_start))
            return target_dir

        all_new_rows = pa.concat_tables(all_new_tables, promote_options="default")
        total_new = len(all_new_rows)
        logger.info("incremental_update_new_rows_exported", count=total_new, new_start=str(new_start))

        # 2. Write new rows to enrichment input dir (all as test split, empty train/valid)
        enrich_input_dir = stage_root / "00_new_base"
        enrich_input_dir.mkdir()
        empty = all_new_rows.schema.empty_table()
        pq.write_table(empty, enrich_input_dir / "train.parquet")
        pq.write_table(empty, enrich_input_dir / "valid.parquet")
        pq.write_table(all_new_rows, enrich_input_dir / "test.parquet")
        # Copy manifest from export
        manifest_src = new_export_dir / "manifest.json"
        if manifest_src.exists():
            shutil.copy2(manifest_src, enrich_input_dir / "manifest.json")

        # 3. Run enrichment steps on new rows only
        steps = [s for s in list(spec.get("enrichment_steps") or []) if s != "base_export"]
        current_dir = enrich_input_dir
        for idx, step in enumerate(steps, start=1):
            is_last = idx == len(steps)
            next_dir = stage_root / (f"enrich_{canonical_snapshot_id}" if is_last else f"enrich_{idx:02d}_{step}")
            _run_enrichment_step(step, current_dir, next_dir)
            current_dir = next_dir

        if current_dir.name != f"enrich_{canonical_snapshot_id}":
            final_enrich_dir = stage_root / f"enrich_{canonical_snapshot_id}"
            shutil.copytree(current_dir, final_enrich_dir)
            current_dir = final_enrich_dir

        # 4. Materialize runtime backfills for new rows
        await _materialize_runtime_backfills(current_dir)

        # 5. Load enriched new test rows
        new_test = pq.read_table(str(current_dir / "test.parquet"))
        logger.info("incremental_update_new_rows_enriched", count=len(new_test))

        # 6. Load existing canonical splits
        existing_train = pq.read_table(str(target_dir / "train.parquet"))
        existing_valid = pq.read_table(str(target_dir / "valid.parquet"))
        existing_test = pq.read_table(str(target_dir / "test.parquet"))

        # 7. Append new rows to test split; align schemas (new columns get null in old rows)
        new_test = _coerce_schema(new_test, existing_test.schema)
        merged_test = pa.concat_tables([existing_test, new_test], promote_options="default")
        if "event_date" in merged_test.column_names:
            sort_idx = pc.sort_indices(merged_test, sort_keys=[("event_date", "ascending")])
            merged_test = merged_test.take(sort_idx)

        # 8. Write merged canonical to staging
        merged_dir = stage_root / canonical_snapshot_id
        merged_dir.mkdir()
        pq.write_table(existing_train, str(merged_dir / "train.parquet"))
        pq.write_table(existing_valid, str(merged_dir / "valid.parquet"))
        pq.write_table(merged_test, str(merged_dir / "test.parquet"))

        # 9. Write manifest and atomic swap
        _write_canonical_manifest(merged_dir, canonical_snapshot_id, spec)
        _atomic_swap_directory(merged_dir, target_dir)

    logger.info(
        "incremental_update_canonical_done",
        canonical_snapshot_id=canonical_snapshot_id,
        new_rows_added=total_new,
    )
    return target_dir


async def materialize_canonical_snapshot(
    requested_snapshot_id: str,
) -> Path:
    """Persist runtime-backfilled features into an existing canonical snapshot."""
    resolution = resolve_snapshot(requested_snapshot_id, allow_legacy_passthrough=False)
    canonical_snapshot_id = resolution.canonical_snapshot_id
    spec = get_canonical_spec(canonical_snapshot_id)
    target_dir = canonical_snapshot_dir(canonical_snapshot_id)
    if not target_dir.exists():
        raise FileNotFoundError(f"Canonical snapshot not found: {target_dir}")

    with tempfile.TemporaryDirectory(
        prefix=f".{canonical_snapshot_id}.materialize.",
        dir=str(target_dir.parent),
    ) as stage_root_str:
        stage_root = Path(stage_root_str)
        current_dir = stage_root / canonical_snapshot_id
        shutil.copytree(target_dir, current_dir)
        materialized_columns = await _materialize_runtime_backfills(current_dir)
        _write_canonical_manifest(
            current_dir,
            canonical_snapshot_id,
            spec,
            extra_fields={
                "materialized_feature_columns": materialized_columns,
                "materialization_mode": "runtime_backfill_persisted_v1",
            },
        )
        _atomic_swap_directory(current_dir, target_dir)

    return target_dir