#!/usr/bin/env python3
"""
Mock HIBERNATE test harness.

Fires posture=HIBERNATE into HZ, then watches CH trade_events for the
resulting exit reason.  Restores posture on exit.

Usage:
    python3 mock_hibernate_test.py [--timeout 120] [--restore APEX]

SAFE: does not touch the engine or any live process.  It only writes
to DOLPHIN_SAFETY (posture), which is the same channel MHS uses.
BLUE will see the HIBERNATE posture on its next _read_posture() call
(cached 10 s) and arm protect mode if a position is open.
"""
import sys, json, time, signal, argparse
from datetime import datetime, timezone
sys.path.insert(0, '/mnt/dolphinng5_predict/prod')

HZ_CLUSTER = "dolphin"
HZ_HOST    = "127.0.0.1:5701"
CH_URL     = "http://localhost:8123/"
CH_AUTH    = ("dolphin", "dolphin_ch_2026")


# ── helpers ──────────────────────────────────────────────────────────────────

def now_iso():
    return datetime.now(timezone.utc).isoformat()

def ch_query(sql: str) -> str:
    import urllib.request, urllib.parse
    data = sql.encode()
    req = urllib.request.Request(CH_URL, data=data,
          headers={"Authorization": "Basic " +
                   __import__('base64').b64encode(b"dolphin:dolphin_ch_2026").decode()})
    with urllib.request.urlopen(req, timeout=10) as r:
        return r.read().decode().strip()

def get_latest_trade_ts() -> str:
    """Return the ts of the most recent trade_events row."""
    return ch_query("SELECT max(ts) FROM dolphin.trade_events FORMAT TabSeparated")

def poll_new_exits(since_ts: str, timeout_s: int):
    """Poll CH every 3 s for new exit rows after since_ts; return list of rows."""
    deadline = time.time() + timeout_s
    while time.time() < deadline:
        time.sleep(3)
        rows = ch_query(
            f"SELECT ts, asset, exit_reason, pnl FROM dolphin.trade_events "
            f"WHERE ts > '{since_ts}' ORDER BY ts FORMAT TabSeparated"
        )
        if rows:
            return rows
    return None


# ── main ─────────────────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--timeout', type=int, default=120,
                    help='Seconds to wait for an exit before restoring posture (default 120)')
    ap.add_argument('--restore', default='APEX',
                    help='Posture to restore after test (default APEX)')
    args = ap.parse_args()

    import hazelcast
    hz = hazelcast.HazelcastClient(cluster_name=HZ_CLUSTER,
                                   cluster_members=[HZ_HOST])
    safety_map = hz.get_map("DOLPHIN_SAFETY").blocking()

    # Read current posture so we can restore it
    raw = safety_map.get("latest")
    if raw:
        try:
            current = json.loads(raw) if isinstance(raw, str) else raw
            original_posture = current.get("posture", "APEX")
        except Exception:
            original_posture = raw if isinstance(raw, str) else "APEX"
    else:
        original_posture = "APEX"

    print(f"[{now_iso()}] Current posture: {original_posture}")

    # Check for open position
    open_pos = ch_query(
        "SELECT asset FROM dolphin.trade_events "
        "GROUP BY asset HAVING count() > 0 "
        "LIMIT 1 FORMAT TabSeparated"
    )
    print(f"[{now_iso()}] Open position check (recent asset): {open_pos or 'none detected via CH'}")

    # Snapshot current latest trade ts
    snap_ts = get_latest_trade_ts()
    print(f"[{now_iso()}] Snapshot trade ts: {snap_ts}")

    # Write HIBERNATE posture
    hibernate_payload = json.dumps({
        "posture": "HIBERNATE",
        "source": "mock_hibernate_test",
        "ts": now_iso(),
    })

    def restore():
        restore_payload = json.dumps({
            "posture": args.restore,
            "source": "mock_hibernate_test_restore",
            "ts": now_iso(),
        })
        safety_map.put("latest", restore_payload)
        safety_map.put("posture", args.restore)
        print(f"\n[{now_iso()}] ✅ Posture restored → {args.restore}")
        hz.shutdown()

    # Trap Ctrl-C to always restore
    def _sig(s, f):
        print(f"\n[{now_iso()}] Interrupted — restoring posture...")
        restore()
        sys.exit(0)
    signal.signal(signal.SIGINT,  _sig)
    signal.signal(signal.SIGTERM, _sig)

    print(f"\n[{now_iso()}] 🔥 Firing HIBERNATE → DOLPHIN_SAFETY.latest")
    safety_map.put("latest", hibernate_payload)
    safety_map.put("posture", "HIBERNATE")
    print(f"[{now_iso()}] Waiting up to {args.timeout}s for exit "
          f"(BLUE reads posture every 10 s)...\n")

    result = poll_new_exits(snap_ts, args.timeout)

    if result:
        print(f"[{now_iso()}] 🎯 Exit detected:")
        for line in result.strip().split('\n'):
            cols = line.split('\t')
            if len(cols) >= 4:
                print(f"  ts={cols[0]}  asset={cols[1]}  reason={cols[2]}  pnl={cols[3]}")
            else:
                print(f"  {line}")
        # Evaluate result
        reasons = [l.split('\t')[2] for l in result.strip().split('\n') if l]
        if any('HIBERNATE_TP' in r for r in reasons):
            print("\n✅ PASS: HIBERNATE_TP fired — protect mode worked, position exited at TP")
        elif any('HIBERNATE_SL' in r for r in reasons):
            print("\n✅ PASS: HIBERNATE_SL fired — protect mode worked, SL guard triggered")
        elif any('HIBERNATE_MAXHOLD' in r for r in reasons):
            print("\n✅ PASS: HIBERNATE_MAXHOLD fired — protect mode held to MAX_HOLD")
        elif any('HIBERNATE_HALT' in r for r in reasons):
            print("\n❌ FAIL: HIBERNATE_HALT fired — protect mode NOT active (was position open?)")
        else:
            print(f"\n⚠️  UNKNOWN exit reason(s): {reasons}")
    else:
        print(f"[{now_iso()}] ⏰ Timeout — no exit detected in {args.timeout}s")
        print("  Possible causes: no open position when HIBERNATE fired, or BLUE "
              "not running / posture cache still active.")

    restore()


if __name__ == '__main__':
    main()