#!/usr/bin/env python3
"""
scripts/test_sendpath.py — prove the send path is inside a slot. Or that it isn't.

    venv/bin/python scripts/test_sendpath.py

────────────────────────────────────────────────────────────────────────────
WHY THIS EXISTS, AND HOW IT DIFFERS FROM /bench
────────────────────────────────────────────────────────────────────────────
Operator: "how you improve the speed 400 ms — make test show us."

Fair, and the honest answer was that /pipeline could not show it. A
pipeline trace is only written when a real signal reaches the executor,
and nothing has reached it since the fix — so /pipeline still displays
three attempts from four days ago at 1.19s, measured on code that no
longer exists. A screen that cannot be refreshed cannot be evidence.

/bench does not answer it either. /bench times each stage in ISOLATION on
a fresh connection: it tells you the network is fine, which was never the
question. The question is whether the CACHES work, and an isolated timing
deliberately avoids them.

So this walks the real thing. Every function called below is the exact
function solana_executor calls on a live send — `_cached_mint_safety`,
`_cached_priority_fee`, `_cached_reserve`, `_cached_alt_accounts`,
`_fresh_blockhash` — through the same hot_state cache, with the same keys.

Each is timed TWICE:

    COLD   the first call, cache empty, the network round trip
    WARM   the second call, which is what a real trade actually pays

The gap between those two columns IS the fix. `freeze` was 758 ms and
51% of the send path; if the warm column does not show it near zero, the
warm cache is not working and this prints that in as many words.

────────────────────────────────────────────────────────────────────────────
WHAT IT NEVER DOES
────────────────────────────────────────────────────────────────────────────
No key is loaded. No transaction is built, signed, simulated or sent. No
bundle reaches Jito. It reads public state over RPC and times the reads.
Safe to run while the bot is live — it shares the same warm cache, so it
leaves the caches hotter than it found them.
"""

from __future__ import annotations

import argparse
import asyncio
import os
import sys
import time
from pathlib import Path
from typing import Any, Callable, Optional

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

from modules.env_file import load_env_file

load_env_file()

_G, _Y, _R, _B, _0 = "\033[32m", "\033[33m", "\033[31m", "\033[1m", "\033[0m"
_SLOT_MS = 400.0


async def _time(fn: Callable[[], Any]) -> tuple[Optional[float], str]:
    """One call, in milliseconds. A failure returns None, never a zero.

    A stage that errors and reports 0 ms is the single most misleading
    number this script could print — it would look like the fastest stage
    on the page.
    """
    t0 = time.perf_counter()
    try:
        result = fn()
        if asyncio.iscoroutine(result):
            await result
    except Exception as exc:  # noqa: BLE001
        return None, str(exc)[:70]
    return (time.perf_counter() - t0) * 1000.0, ""


async def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--warm-runs", type=int, default=3,
                    help="warm samples per stage (default 3, the median is used)")
    args = ap.parse_args()

    from constants import SOLANA_TOKENS
    from modules.solana_executor import (
        _cached_alt_accounts, _cached_mint_safety, _cached_priority_fee,
        _cached_reserve, _fresh_blockhash, get_shared_client,
    )

    rpc = (os.getenv("SOLANA_RPC_PRIMARY") or os.getenv("SOLANA_RPC_URL") or "").strip()
    if not rpc:
        print("\n  no RPC configured (SOLANA_RPC_PRIMARY) — nothing to measure\n")
        return 1

    client = get_shared_client()
    sol, usdc = SOLANA_TOKENS["SOL"], SOLANA_TOKENS["USDC"]

    # Exactly the calls a send makes, with the same cache keys.
    stages: list[tuple[str, Callable[[], Any], bool]] = [
        ("freeze SOL", lambda: _cached_mint_safety(client, rpc, "SOL", sol["mint"]), True),
        ("freeze USDC", lambda: _cached_mint_safety(client, rpc, "USDC", usdc["mint"]), True),
        ("priority fee", lambda: _cached_priority_fee(client, rpc), True),
        ("reserve", lambda: _cached_reserve(client, "USDC", usdc["mint"], rpc), True),
        ("alt tables", lambda: _cached_alt_accounts(client, rpc, []), True),
        # Not cacheable in the same sense — a blockhash has a real
        # expiry — but hot_state serves it behind a short TTL, so the
        # caller should not be paying the round trip either.
        ("blockhash", lambda: _fresh_blockhash(client, rpc), False),
    ]

    print()
    print(f"{_B}  ⏱  send-path test — the real cached calls, nothing signed{_0}")
    print("     COLD = cache empty · WARM = what a live trade actually pays")
    print("  " + "─" * 64)
    print(f"  {'stage':<16}{'COLD':>10}{'WARM':>10}   saved")
    print("  " + "─" * 64)

    warm_total = 0.0
    cold_total = 0.0
    broken: list[str] = []

    for name, fn, cacheable in stages:
        cold, err = await _time(fn)
        if cold is None:
            broken.append(f"{name}: {err}")
            print(f"  {name:<16}{_R}{'failed':>10}{_0}   {err[:30]}")
            continue

        warms: list[float] = []
        for _ in range(max(1, args.warm_runs)):
            w, werr = await _time(fn)
            if w is not None:
                warms.append(w)
        if not warms:
            broken.append(f"{name}: warm call failed")
            continue
        warm = sorted(warms)[len(warms) // 2]

        cold_total += cold
        warm_total += warm
        saved = cold - warm
        # Colour the WARM column, because that is the number a trade pays.
        colour = _G if warm < 20 else (_Y if warm < 100 else _R)
        flag = "" if cacheable else "  (short TTL by design)"
        print(f"  {name:<16}{cold:>9.0f}ms{colour}{warm:>9.0f}ms{_0}"
              f"   {saved:>6.0f}ms{flag}")

    print("  " + "─" * 64)
    print(f"  {'TOTAL':<16}{cold_total:>9.0f}ms{warm_total:>9.0f}ms"
          f"   {cold_total - warm_total:>6.0f}ms")
    print()

    if broken:
        print(f"  {_Y}Some stages could not be measured:{_0}")
        for b in broken:
            print(f"    {b}")
        print("  The totals above exclude them, so they understate the real")
        print("  pipeline. Fix these before reading this as a pass.")
        print()
        return 1

    # A live send also pays two quotes and two swap-instruction builds.
    # Those are NOT cacheable — they are the market — so /bench is the
    # right measure for them and its figure is used here rather than
    # re-measuring it badly.
    quotes_ms = float(os.getenv("SENDPATH_QUOTE_BUDGET_MS", "120"))
    projected = warm_total + quotes_ms

    print(f"  cached state, warm      {warm_total:>7.0f} ms")
    print(f"  + quotes and swap-ix    {quotes_ms:>7.0f} ms   "
          f"(uncacheable — /bench measures these)")
    print(f"  {_B}= signal to signed      {projected:>7.0f} ms{_0}")
    print(f"    a Solana slot is        {_SLOT_MS:>7.0f} ms")
    print()

    if projected <= _SLOT_MS:
        print(f"  {_G}✓ {projected:.0f} ms — inside one slot.{_0} The warm cache is")
        print("    working. Latency is not what stops a trade; /observe for")
        print("    whether the round trip is positive at all.")
    else:
        print(f"  {_R}✗ {projected:.0f} ms — longer than a slot.{_0}")
        print("    The warm column above names which stage is still paying a")
        print("    round trip. A cached stage over 100 ms warm means the cache")
        print("    is missing — check hot_state's TTL for it and that")
        print("    warm_mint_safety ran at startup.")
    print()

    # The freeze stage specifically, because it was 51% of the pipeline and
    # is the thing the operator asked to see proven.
    print(f"  {_B}The one that mattered:{_0} `freeze` was 758 ms p50 and 51% of the")
    print("  send path. Its WARM column above is what a trade pays now.")
    print()
    return 0


if __name__ == "__main__":
    raise SystemExit(asyncio.run(main()))