#!/usr/bin/env python
"""Seed the signal store: refresh prices, run inference, build comparison tables.

Runnable locally, in Colab, or from the Space's ZeroGPU function. Progress is
checkpointed after every batch, so an interrupted run resumes where it stopped
rather than paying for the same inference twice.

    python scripts/seed_store.py --plan v1 --dry-run
    python scripts/seed_store.py --plan v1 --prices-only
    python scripts/seed_store.py --plan v1 --push

Nothing is recomputed that the manifest already covers.
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import sys
import time
from dataclasses import dataclass, asdict, field
from pathlib import Path

import numpy as np
import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from src import catalog, comparisons, config  # noqa: E402
from src.adapters import build_windows, get_adapter  # noqa: E402
from src.data import refresh  # noqa: E402
from src.store import SignalStore  # noqa: E402

log = logging.getLogger("seed")


# --------------------------------------------------------------------------
# Seed plans
# --------------------------------------------------------------------------


@dataclass(frozen=True)
class SeedTarget:
    model_slug: str
    asset: str
    timeframe: str
    years: float
    # Placeholder targets are written with inference_version PLACEHOLDER and
    # are replaced the moment a real run covers the same slice.
    placeholder: bool = False

    @property
    def key(self) -> str:
        return f"{self.model_slug}|{self.asset}|{self.timeframe}"


CRYPTO = ["BTC-USD", "ETH-USD", "SOL-USD"]
EQUITIES = ["SPY", "QQQ", "NVDA"]


def plan_v1() -> list[SeedTarget]:
    """The v1 seed: every seedable model on daily bars for the whole universe,
    plus intraday coverage for the fast Chronos-Bolt family on crypto.

    Daily is the comparison backbone -- every model sees exactly the same bars
    on every asset, so leaderboard differences are the model, not the coverage.
    Intraday is added where inference is cheap enough to be honest about.
    """
    targets: list[SeedTarget] = []
    all_assets = CRYPTO + EQUITIES

    # Backbone: every model x every asset, daily.
    for model in config.SEEDABLE_MODELS:
        for asset in all_assets:
            targets.append(SeedTarget(model, asset, "1d", 3.0))

    # Intraday: the bolt family plus baselines on crypto.
    intraday_models = [m for m in config.SEEDABLE_MODELS
                       if m.startswith("chronos-bolt") or m.startswith("baseline")]
    for model in intraday_models:
        for asset in CRYPTO:
            targets.append(SeedTarget(model, asset, "1h", 1.0))

    # 15-minute: the small model and the naive baseline, crypto only.
    for model in ("chronos-bolt-small", "baseline-naive"):
        for asset in CRYPTO:
            targets.append(SeedTarget(model, asset, "15m", 0.25))

    # Hourly equities for the reference model, capped by provider depth.
    for asset in EQUITIES:
        targets.append(SeedTarget("chronos-bolt-small", asset, "1h", 1.5))

    return targets


def plan_smoke() -> list[SeedTarget]:
    """One model, one asset, six months -- proves the pipeline end to end."""
    return [SeedTarget("chronos-bolt-small", "BTC-USD", "1d", 0.5)]


PLANS = {"v1": plan_v1, "smoke": plan_smoke}


# --------------------------------------------------------------------------
# Checkpointing
# --------------------------------------------------------------------------


@dataclass
class Checkpoint:
    path: Path
    done: dict[str, str] = field(default_factory=dict)  # key -> last ts written
    failed: dict[str, str] = field(default_factory=dict)

    @classmethod
    def load(cls, path: str | os.PathLike) -> "Checkpoint":
        p = Path(path)
        if p.exists():
            try:
                raw = json.loads(p.read_text())
                return cls(path=p, done=raw.get("done", {}), failed=raw.get("failed", {}))
            except json.JSONDecodeError:
                log.warning("checkpoint at %s was corrupt; starting fresh", p)
        return cls(path=p)

    def save(self) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.path.write_text(json.dumps(
            {"done": self.done, "failed": self.failed,
             "updated_at": pd.Timestamp.now(tz="UTC").isoformat()}, indent=2))

    def mark(self, key: str, last_ts) -> None:
        self.done[key] = str(last_ts)
        self.failed.pop(key, None)
        self.save()

    def mark_failed(self, key: str, reason: str) -> None:
        self.failed[key] = reason
        self.save()

    def last_ts(self, key: str) -> pd.Timestamp | None:
        v = self.done.get(key)
        return pd.Timestamp(v) if v else None


# --------------------------------------------------------------------------
# Steps
# --------------------------------------------------------------------------


def refresh_prices(store: SignalStore, targets: list[SeedTarget]) -> list[str]:
    """Fetch only the price ranges the store is missing."""
    notes = []
    wanted: dict[tuple[str, str], float] = {}
    for t in targets:
        k = (t.asset, t.timeframe)
        wanted[k] = max(wanted.get(k, 0.0), t.years)

    end = pd.Timestamp.now(tz="UTC").floor("h")
    for (asset, tf), years in sorted(wanted.items()):
        start = end - pd.Timedelta(days=int(365 * years) + 30)
        rep = refresh(store, asset, tf, start, end)
        log.info("%s", rep.summary())
        notes.append(rep.summary())
        for n in rep.boundary_notes:
            log.info("  boundary: %s", n)
            notes.append(f"  boundary: {n}")
    return notes


def seed_target(store: SignalStore, target: SeedTarget, ckpt: Checkpoint,
                *, batch_size: int = 256, device: str | None = None,
                force_placeholder: bool = False) -> str:
    """Run inference for one (model, asset, timeframe) and write the slice."""
    spec = config.SEED_MODELS.get(target.model_slug)
    if spec is None:
        return f"SKIP {target.key}: unknown model"

    prices = store.get_prices(target.asset, target.timeframe)
    if prices.empty:
        return f"SKIP {target.key}: no price coverage"

    end = prices.index[-1]
    start = end - pd.Timedelta(days=int(365 * target.years))
    prices = prices[prices.index >= start]
    close = prices["close"]

    use_placeholder = target.placeholder or force_placeholder
    family = "placeholder" if use_placeholder else spec.family
    ctx_len = min(spec.context_len, max(64, len(close) // 3))

    adapter = get_adapter(family, spec.model_id, context_len=ctx_len, device=device)
    adapter.load()
    revision = adapter.resolved_revision
    version = adapter.inference_version()

    stamps, windows = build_windows(close, ctx_len)
    if len(stamps) == 0:
        return f"SKIP {target.key}: only {len(close)} bars, need > {ctx_len}"

    # Idempotency: never recompute what the manifest already covers. The range
    # to check is the one the windows actually produce -- signals start a full
    # context window after the first price bar, so checking the price range
    # would always report the leading context as an uncovered gap.
    missing = store.missing_ranges(target.model_slug, revision, target.asset,
                                   target.timeframe, stamps[0], stamps[-1])
    if not missing:
        return f"SKIP {target.key}: already covered by the manifest"

    resume_from = ckpt.last_ts(target.key)
    if resume_from is not None:
        keep = stamps > resume_from
        stamps, windows = stamps[keep], windows[keep]
        if len(stamps) == 0:
            return f"SKIP {target.key}: checkpoint says complete"

    t0 = time.perf_counter()
    written = 0
    for i in range(0, len(stamps), batch_size):
        bs, bw = stamps[i:i + batch_size], windows[i:i + batch_size]
        forecast = adapter.predict(bw)
        frame = forecast.as_frame(bs, version)
        store.write_signals(
            target.model_slug, spec.model_id, revision, target.asset,
            target.timeframe, frame,
            inference_version=version, contributed_by="seed",
        )
        written += len(frame)
        ckpt.mark(target.key, bs[-1])
        log.info("  %s %d/%d", target.key, min(i + batch_size, len(stamps)), len(stamps))

    dt = time.perf_counter() - t0
    tag = " [PLACEHOLDER]" if use_placeholder else ""
    return (f"OK   {target.key}: {written} steps in {dt:.1f}s "
            f"({dt / max(written, 1) * 1000:.0f} ms/step){tag}")


# --------------------------------------------------------------------------
# Main
# --------------------------------------------------------------------------


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description="Seed the bit signal store")
    ap.add_argument("--plan", default="v1", choices=sorted(PLANS))
    ap.add_argument("--store-root", default=".cache/store")
    ap.add_argument("--checkpoint", default=".cache/seed_checkpoint.json")
    ap.add_argument("--repo", default=config.STORE_REPO)
    ap.add_argument("--batch-size", type=int, default=256)
    ap.add_argument("--device", default=None)
    ap.add_argument("--prices-only", action="store_true")
    ap.add_argument("--skip-prices", action="store_true")
    ap.add_argument("--placeholder-only", action="store_true",
                    help="write labelled synthetic signals instead of running models")
    ap.add_argument("--no-comparisons", action="store_true")
    ap.add_argument("--push", action="store_true", help="commit to the Hub when done")
    ap.add_argument("--offline", action="store_true")
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--only", default=None, help="substring filter on target keys")
    args = ap.parse_args(argv)

    logging.basicConfig(level=logging.INFO, format="%(message)s")

    targets = PLANS[args.plan]()
    if args.only:
        targets = [t for t in targets if args.only in t.key]

    if args.dry_run:
        print(f"plan={args.plan}  targets={len(targets)}")
        for t in targets:
            print(f"  {t.key:48s} years={t.years:<5g} placeholder={t.placeholder}")
        return 0

    store = SignalStore(repo_id=None if args.offline else args.repo,
                        local_root=args.store_root, offline=args.offline)
    ckpt = Checkpoint.load(args.checkpoint)
    results: list[str] = []

    if not args.skip_prices:
        log.info("== refreshing prices ==")
        results.extend(refresh_prices(store, targets))

    if not args.prices_only:
        log.info("== running inference ==")
        for t in targets:
            try:
                msg = seed_target(store, t, ckpt, batch_size=args.batch_size,
                                  device=args.device,
                                  force_placeholder=args.placeholder_only)
            except Exception as e:
                log.exception("target %s failed", t.key)
                ckpt.mark_failed(t.key, f"{type(e).__name__}: {e}")
                msg = f"FAIL {t.key}: {type(e).__name__}: {e}"
            log.info("%s", msg)
            results.append(msg)

    if not args.no_comparisons:
        log.info("== regenerating comparison tables ==")
        report = comparisons.regenerate(store)
        results.append(
            f"comparisons: perf={len(report.model_performance)} "
            f"calib={len(report.calibration)} dir={len(report.directional)} "
            f"heatmap={len(report.heatmap)}"
        )
        log.info("%s", results[-1])

        log.info("== building catalog ==")
        cat = catalog.build(store)
        results.append(cat.summary())
        log.info("%s", cat.summary())

    if args.push and not args.offline:
        log.info("== pushing to %s ==", args.repo)
        oid = store.flush(f"Seed store ({args.plan})")
        log.info("commit: %s", oid)
        results.append(f"pushed commit {oid}")

    print("\n=== SEED SUMMARY ===")
    for r in results:
        print(r)
    failures = [r for r in results if r.startswith("FAIL")]
    return 1 if failures else 0


if __name__ == "__main__":
    raise SystemExit(main())