src/extension.py

"""User-funded coverage extension and add-model, on ZeroGPU.

Inference runs inside a `@spaces.GPU` function so it draws on the *signed-in
user's* ZeroGPU quota, not the Space owner's. Anonymous visitors keep full read
and backtest access; only extension is gated.

Safety properties enforced here:

* **Dedup** — a range the manifest already covers is never recomputed.
* **Caps** — per-request range limits keep one user from monopolising the queue.
* **Single writer** — a process-wide lock serialises commits, so concurrent
  extensions cannot interleave and corrupt the manifest.
* **Allow-list** — only vetted adapter families load, and model ids are
  validated before they reach the Hub.
"""

from __future__ import annotations

import logging
import os
import threading
from dataclasses import dataclass

import gradio as gr
import pandas as pd

from . import catalog, comparisons, config, runtime
from .adapters import AdapterError, ModelNotAllowed, build_windows, get_adapter, validate_model_id
from .store import _utc

log = logging.getLogger("bit.extension")

# One writer at a time. Commits to the store are serialised process-wide so a
# second extension cannot land between another's data write and its manifest
# update.
_WRITE_LOCK = threading.Lock()

try:
    import spaces  # provided by the ZeroGPU runtime
    HAS_SPACES = True
except Exception:  # running locally or on CPU-only hardware
    spaces = None
    HAS_SPACES = False


def _gpu(duration=120):
    """Apply @spaces.GPU when the runtime offers it, otherwise run on CPU."""
    def deco(fn):
        if HAS_SPACES:
            return spaces.GPU(duration=duration)(fn)
        return fn
    return deco


class ExtensionError(RuntimeError):
    pass


class QuotaExhausted(ExtensionError):
    pass


# --------------------------------------------------------------------------
# Estimation & guardrails
# --------------------------------------------------------------------------


@dataclass
class Estimate:
    model_slug: str
    asset: str
    timeframe: str
    start: pd.Timestamp
    end: pd.Timestamp
    steps: int
    already_covered: bool
    capped: bool
    cap_days: int
    note: str = ""


def estimate(model_slug: str, asset: str, timeframe: str, start, end) -> Estimate:
    if model_slug not in config.SEED_MODELS:
        raise ExtensionError(f"unknown model {model_slug!r}")
    if asset not in config.ASSETS:
        raise ExtensionError(f"unknown asset {asset!r}")
    if timeframe not in config.TIMEFRAMES:
        raise ExtensionError(f"unknown timeframe {timeframe!r}")

    try:
        s, e = _utc(start), _utc(end)
    except Exception as exc:
        raise ExtensionError(f"could not parse the date range: {exc}") from exc
    if s >= e:
        raise ExtensionError("start must be before end")

    cap_days = config.CAPS.max_days.get(timeframe, 365)
    capped = (e - s).days > cap_days
    if capped:
        s = e - pd.Timedelta(days=cap_days)

    store = runtime.get_store()
    prices = store.get_prices(asset, timeframe, s, e)
    if prices.empty:
        raise ExtensionError(
            f"No cached prices for {asset} {timeframe} in that range. "
            "Price coverage has to exist before signals can be generated."
        )

    spec = config.SEED_MODELS[model_slug]
    ctx = min(spec.context_len, max(64, len(prices) // 3))
    steps = max(0, len(prices) - ctx)
    if steps > config.CAPS.max_steps_per_run:
        steps = config.CAPS.max_steps_per_run

    # Dedup must compare against the range this request would actually
    # *produce*, not the range the user typed. A forecast needs a full trailing
    # context window, so the first producible timestamp sits `ctx` bars after
    # the start of the price slice. Comparing the typed range instead would
    # report an already-covered slice as uncovered and pay for it twice.
    stamps, _ = build_windows(prices["close"], ctx)
    if len(stamps) == 0:
        produced_start, produced_end = s, e
        steps = 0
    else:
        produced_start = _utc(stamps[0])
        produced_end = _utc(stamps[min(steps, len(stamps)) - 1]) if steps else produced_start

    # Revision is resolved lazily to avoid a Hub round-trip on every keystroke,
    # so this matches any revision of the model.
    covered = any(
        ent.model_slug == model_slug and ent.asset == asset
        and ent.timeframe == timeframe
        and _utc(ent.start_ts) <= produced_start and _utc(ent.end_ts) >= produced_end
        for ent in store.load_manifest().signals.values()
    )

    note = ""
    if capped:
        note = (f"Range trimmed to the {cap_days}-day cap for {timeframe} bars.")
    return Estimate(model_slug, asset, timeframe, s, e, steps, covered,
                    capped, cap_days, note)


# --------------------------------------------------------------------------
# GPU inference
# --------------------------------------------------------------------------


@_gpu(duration=180)
def run_inference(model_id: str, family: str, values: list, ctx_len: int) -> dict:
    """Inference on the caller's ZeroGPU allocation.

    Kept deliberately small and picklable: it takes plain values and returns
    plain lists, so nothing in the store or the app leaks into the GPU worker.
    """
    import numpy as np

    adapter = get_adapter(family, model_id, context_len=ctx_len)
    adapter.load()
    windows = np.asarray(values, dtype="float32")
    forecast = adapter.predict(windows)
    return {
        "q10": forecast.q10.tolist(),
        "q50": forecast.q50.tolist(),
        "q90": forecast.q90.tolist(),
        "context_len": int(forecast.context_len),
        "revision": adapter.resolved_revision,
        "inference_version": adapter.inference_version(),
    }


def _is_quota_error(exc: Exception) -> bool:
    text = f"{type(exc).__name__} {exc}".lower()
    return any(k in text for k in ("quota", "gpu task aborted", "exceeded",
                                   "no gpu available", "zerogpu"))


QUOTA_FALLBACK = (
    '<div class="bit-note bit-note-danger">'
    "Your ZeroGPU quota is exhausted, so this extension could not run. "
    "The quota refills over time. If you need to run a large batch now, "
    f'<a href="https://huggingface.co/spaces/{config.SPACE_REPO}?duplicate=true" '
    'target="_blank" rel="noopener">duplicate this Space</a> and run it on your '
    "own hardware — the signal store is public, so a duplicate reads the same data."
    "</div>"
)


# --------------------------------------------------------------------------
# The extend flow
# --------------------------------------------------------------------------


def extend_coverage(model_slug: str, asset: str, timeframe: str, start, end,
                    username: str = "anonymous", progress=None) -> str:
    """Dedup, estimate, run inference, commit, regenerate comparisons."""
    est = estimate(model_slug, asset, timeframe, start, end)
    if est.already_covered:
        return ('<div class="bit-note">That range is already covered — nothing was '
                "recomputed. Coverage is deduplicated against the manifest.</div>")
    if est.steps <= 0:
        return ('<div class="bit-note bit-note-danger">Not enough cached price '
                "history in that range to build a single context window.</div>")

    spec = config.SEED_MODELS[model_slug]
    store = runtime.get_store()
    prices = store.get_prices(asset, timeframe, est.start, est.end)
    close = prices["close"]
    ctx = min(spec.context_len, max(64, len(prices) // 3))

    if progress:
        progress(0.15, desc=f"Preparing {est.steps} context windows")
    stamps, windows = build_windows(close, ctx)
    if len(stamps) == 0:
        return ('<div class="bit-note bit-note-danger">Not enough bars for a '
                "context window.</div>")
    stamps, windows = stamps[:est.steps], windows[:est.steps]

    if progress:
        progress(0.35, desc=f"Running {spec.display} on your GPU quota")
    try:
        out = run_inference(spec.model_id, spec.family, windows.tolist(), ctx)
    except Exception as exc:
        if _is_quota_error(exc):
            log.warning("ZeroGPU quota exhausted for %s: %s", username, exc)
            return QUOTA_FALLBACK
        log.exception("extension inference failed")
        return (f'<div class="bit-note bit-note-danger">Inference failed: '
                f"{type(exc).__name__}: {exc}</div>")

    frame = pd.DataFrame({
        "ts": stamps, "q10": out["q10"], "q50": out["q50"], "q90": out["q90"],
        "context_len": out["context_len"], "inference_version": out["inference_version"],
    })

    if progress:
        progress(0.75, desc="Committing to the signal store")
    with _WRITE_LOCK:
        entry = store.write_signals(
            model_slug, spec.model_id, out["revision"], asset, timeframe, frame,
            inference_version=out["inference_version"],
            contributed_by=username,
        )
        if progress:
            progress(0.9, desc="Regenerating comparison tables")
        # Coverage just changed, so every derived view is now stale. Both are
        # rebuilt inside the same write lock, before the commit, so readers
        # never see new signals alongside an old leaderboard.
        try:
            comparisons.regenerate(store, assets=[asset])
        except Exception:
            log.exception("comparison regeneration failed (coverage still written)")
        try:
            catalog.build(store)
        except Exception:
            log.exception("catalog rebuild failed (coverage still written)")
        oid = store.flush(f"Extend {model_slug}/{asset}/{timeframe} by @{username}")

    runtime.cache_clear()
    return (f'<div class="bit-note">Coverage extended by <b>@{username}</b> — '
            f"{len(frame):,} new steps for {model_slug} on {asset} {timeframe} "
            f"({entry.start_ts[:10]} → {entry.end_ts[:10]})."
            + (f" Commit <code>{str(oid)[:8]}</code>." if oid else "")
            + "</div>")


def add_model(family: str, model_id: str, username: str = "anonymous") -> str:
    """Smoke-test a user-supplied model, then register it if it passes."""
    try:
        model_id = validate_model_id(model_id)
    except AdapterError as e:
        return f'<div class="bit-note bit-note-danger">{e}</div>'
    if family not in config.ALLOWED_ADAPTER_FAMILIES:
        return (f'<div class="bit-note bit-note-danger">Adapter family {family!r} '
                "is not on the allow-list.</div>")

    store = runtime.get_store()
    asset, timeframe = "BTC-USD", "1d"
    prices = store.get_prices(asset, timeframe)
    if prices.empty:
        return ('<div class="bit-note bit-note-danger">No cached prices to smoke '
                "test against.</div>")

    n = config.CAPS.smoke_test_steps
    close = prices["close"]
    ctx = min(512, max(64, len(close) // 3))
    stamps, windows = build_windows(close, ctx)
    if len(stamps) < n:
        return ('<div class="bit-note bit-note-danger">Not enough history for a '
                f"{n}-step smoke test.</div>")
    stamps, windows = stamps[-n:], windows[-n:]

    try:
        out = run_inference(model_id, family, windows.tolist(), ctx)
    except ModelNotAllowed as e:
        return f'<div class="bit-note bit-note-danger">{e}</div>'
    except Exception as exc:
        if _is_quota_error(exc):
            return QUOTA_FALLBACK
        return (f'<div class="bit-note bit-note-danger">Smoke test failed: '
                f"{type(exc).__name__}: {exc}</div>")

    slug = model_id.split("/")[-1].lower()
    frame = pd.DataFrame({
        "ts": stamps, "q10": out["q10"], "q50": out["q50"], "q90": out["q90"],
        "context_len": out["context_len"], "inference_version": out["inference_version"],
    })
    with _WRITE_LOCK:
        store.write_signals(slug, model_id, out["revision"], asset, timeframe, frame,
                            inference_version=out["inference_version"],
                            contributed_by=username)
        store.flush(f"Add model {model_id} (smoke test) by @{username}")
    runtime.cache_clear()

    return (f'<div class="bit-note">Smoke test passed: <b>{model_id}</b> '
            f"(revision <code>{str(out['revision'])[:8]}</code>) produced {n} "
            f"schema-valid steps and now appears in the coverage map as "
            f"<code>{slug}</code>.</div>")


# --------------------------------------------------------------------------
# Gradio bindings
# --------------------------------------------------------------------------


def status_html() -> str:
    if HAS_SPACES:
        return ('<div class="bit-micro">ZEROGPU AVAILABLE · EXTENSION RUNS ON '
                "YOUR OWN QUOTA WHEN SIGNED IN</div>")
    return ('<div class="bit-note">This Space is running on CPU, so coverage '
            "extension is disabled. Reading and backtesting the existing store "
            "works normally.</div>")


def _username(profile) -> str | None:
    if profile is None:
        return None
    return getattr(profile, "username", None) or getattr(profile, "name", None)


def estimate_ui(model_slug, asset, timeframe, start, end):
    if not (model_slug and asset and timeframe and start and end):
        return '<div class="bit-micro">PICK A MODEL, ASSET, TIMEFRAME AND RANGE</div>'
    try:
        est = estimate(model_slug, asset, timeframe, start, end)
    except ExtensionError as e:
        return f'<div class="bit-note bit-note-danger">{e}</div>'
    if est.already_covered:
        return ('<div class="bit-note">Already covered — running this would '
                "recompute nothing. Pick a wider range.</div>")
    return (f'<div class="bit-note">About <b>{est.steps:,}</b> inference steps for '
            f"{est.start.date()} → {est.end.date()}. "
            + (est.note + " " if est.note else "")
            + "This runs on your ZeroGPU quota once you sign in.</div>")


def extend_ui(model_slug, asset, timeframe, start, end,
              profile: gr.OAuthProfile | None = None, progress=None):
    user = _username(profile)
    if user is None:
        return ('<div class="bit-note bit-note-danger">Sign in with Hugging Face to '
                "extend coverage. Reading and backtesting stay open to everyone; "
                "extension spends your own GPU quota, so it needs an account.</div>",
                runtime.coverage_frame())
    if not HAS_SPACES:
        return (status_html(), runtime.coverage_frame())
    try:
        html = extend_coverage(model_slug, asset, timeframe, start, end,
                               username=user, progress=progress)
    except ExtensionError as e:
        html = f'<div class="bit-note bit-note-danger">{e}</div>'
    except Exception as e:
        log.exception("extend failed")
        html = f'<div class="bit-note bit-note-danger">{type(e).__name__}: {e}</div>'
    return html, runtime.coverage_frame()


def add_model_ui(family, model_id, profile: gr.OAuthProfile | None = None):
    user = _username(profile)
    if user is None:
        return ('<div class="bit-note bit-note-danger">Sign in with Hugging Face to '
                "add a model — the smoke test runs on your GPU quota.</div>",
                runtime.coverage_frame())
    if not HAS_SPACES:
        return (status_html(), runtime.coverage_frame())
    try:
        html = add_model(family, model_id, username=user)
    except Exception as e:
        log.exception("add model failed")
        html = f'<div class="bit-note bit-note-danger">{type(e).__name__}: {e}</div>'
    return html, runtime.coverage_frame()