#!/usr/bin/env python3
"""
Ad-hoc mesh stack exercises (no second Mac, no live LLM on peers).

Run from repo root:
  ~/.mugi/venv/bin/python3 scripts/mesh_stack_smoke.py

Maps loosely to the four proposed stress themes:
  1) Concurrency — many pending task waits per peer (hits configured caps).
  2) Deep context — large task_result text vs pending_result_max_chars.
  3) Disruption — cancel pending when session closes.
  4) Shared-state conflict — not a mesh primitive (notes only).

Exits 0 if all checks pass, 1 if any failed.
"""
from __future__ import annotations

import os
import sys
import threading
import time
import uuid

# Repo root on path
_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _ROOT not in sys.path:
    sys.path.insert(0, _ROOT)

import requests
import uvicorn
from fastapi import FastAPI

from mesh import MeshManager
from mesh.identity import public_key_b64_to_fingerprint
from mesh.protocol import MeshMessage
from mesh.routes import build_mesh_router


def _cfg(**overrides: object) -> dict:
    base = {
        "pending_result_max_total": 20,
        "pending_result_max_per_peer": 5,
        "pending_result_max_chars": 4000,
        "message_max_bytes": 131072,
        "mesh_llm_max_seconds": 120,
    }
    base.update(overrides)
    return base


def _patch_config(m: MeshManager, **kwargs: object) -> None:
    cfg = _cfg(**kwargs)
    m._get_mesh_config = lambda: cfg  # type: ignore[method-assign]


def _fail(tracker: list[str], name: str, detail: str) -> None:
    tracker.append(f"{name}: {detail}")
    print(f"FAIL {name}: {detail}")


def _ok(name: str, detail: str = "") -> None:
    print(f"OK   {name}" + (f" — {detail}" if detail else ""))


def test_pending_per_peer_cap(failures: list[str]) -> None:
    """Theme 1: more concurrent waits than pending_result_max_per_peer."""
    m = MeshManager(port=9201)
    _patch_config(m, pending_result_max_per_peer=5)
    try:
        fp = "fp_storm"
        sid = "sess_storm"
        keys: list[str] = []
        for i in range(5):
            out = m.register_pending_result(
                conversation_id=f"c{i}",
                session_id=sid,
                peer_fingerprint=fp,
                timeout_seconds=60,
            )
            if not out.get("ok"):
                _fail(failures, "pending_cap", f"expected 5 registrations, failed at {i}: {out}")
                return
            keys.append(str(out["key"]))
        sixth = m.register_pending_result(
            conversation_id="c_overflow",
            session_id=sid,
            peer_fingerprint=fp,
            timeout_seconds=60,
        )
        if sixth.get("ok"):
            _fail(failures, "pending_cap", "6th registration should be rejected for same peer")
            return
        err = str(sixth.get("error") or "")
        if "too many" not in err.lower():
            _fail(failures, "pending_cap", f"unexpected error: {sixth}")
            return
        _ok("pending_per_peer_cap", "6th wait rejected as expected")
    finally:
        m.stop()


def test_concurrent_register(failures: list[str]) -> None:
    """Theme 1: thread-safe registration up to per-peer cap."""
    m = MeshManager(port=9202)
    _patch_config(m, pending_result_max_per_peer=10, pending_result_max_total=50)
    try:
        fp = "fp_parallel"
        sid = "sess_parallel"
        results: list[dict] = []
        lock = threading.Lock()

        def job(i: int) -> None:
            out = m.register_pending_result(
                conversation_id=f"p{i}_{uuid.uuid4().hex[:6]}",
                session_id=sid,
                peer_fingerprint=fp,
                timeout_seconds=30,
            )
            with lock:
                results.append(out)

        threads = [threading.Thread(target=job, args=(i,)) for i in range(8)]
        for t in threads:
            t.start()
        for t in threads:
            t.join(timeout=10)
        oks = sum(1 for r in results if r.get("ok"))
        if oks != 8:
            _fail(failures, "concurrent_register", f"expected 8 ok, got {oks}: {results}")
            return
        _ok("concurrent_register", "8 parallel registers for one peer")
    finally:
        m.stop()


def test_large_task_result_truncation(failures: list[str]) -> None:
    """Theme 2: inbound task_result text is capped."""
    m = MeshManager(port=9203)
    max_chars = 3000
    _patch_config(m, pending_result_max_chars=max_chars)
    try:
        fp = "fp_big"
        sid = "sess_big"
        convo = "convo_big"
        reg = m.register_pending_result(
            conversation_id=convo,
            session_id=sid,
            peer_fingerprint=fp,
            timeout_seconds=5,
        )
        if not reg.get("ok"):
            _fail(failures, "truncation", f"register failed: {reg}")
            return
        key = str(reg["key"])
        huge = "B" * (max_chars + 5000)
        msg = MeshMessage(
            type="task_result",
            conversation_id=convo,
            conversation_type="task",
            message_number=2,
            payload={"ok": True, "text": huge},
        )
        m._resolve_pending_result(session_id=sid, peer_fingerprint=fp, msg=msg)
        waited = m.wait_for_pending_result(key=key, timeout_seconds=2, stop_check=lambda: False)
        if not waited.get("ok"):
            _fail(failures, "truncation", f"wait failed: {waited}")
            return
        text = str(waited.get("text") or "")
        if "[Truncated]" not in text:
            _fail(failures, "truncation", "expected [Truncated] suffix on oversized result")
            return
        if len(text) > max_chars + 80:
            _fail(failures, "truncation", f"text longer than cap+margin: {len(text)}")
            return
        _ok("large_task_result_truncation", f"len={len(text)}")
    finally:
        m.stop()


def test_session_cancel_pending(failures: list[str]) -> None:
    """Theme 3: closing session cancels waiters."""
    m = MeshManager(port=9204)
    _patch_config(m)
    try:
        fp = "fp_disc"
        sid = "sess_disc"
        reg = m.register_pending_result(
            conversation_id="c_disc",
            session_id=sid,
            peer_fingerprint=fp,
            timeout_seconds=120,
        )
        key = str(reg["key"])
        results: list = []

        def waiter() -> None:
            results.append(m.wait_for_pending_result(key=key, timeout_seconds=5, stop_check=lambda: False))

        th = threading.Thread(target=waiter, daemon=True)
        th.start()
        time.sleep(0.05)
        n = m.cancel_pending_for_session(sid, "session closed")
        if n < 1:
            _fail(failures, "cancel_pending", f"expected cancel count >=1, got {n}")
            return
        th.join(timeout=3)
        if not results:
            _fail(failures, "cancel_pending", "waiter did not return")
            return
        if results[0].get("error") != "session closed":
            _fail(failures, "cancel_pending", f"unexpected wait result: {results[0]}")
            return
        _ok("session_cancel_pending")
    finally:
        m.stop()


def _wait_listen(host: str, port: int, attempts: int = 40) -> bool:
    import socket

    for _ in range(attempts):
        try:
            with socket.create_connection((host, port), timeout=0.15):
                return True
        except OSError:
            time.sleep(0.05)
    return False


def test_live_handshake_connect(failures: list[str]) -> None:
    """Real TCP: two uvicorn stacks, outbound connect after pending allow."""
    port_a, port_b = 19201, 19202
    ma = MeshManager(port=port_a)
    mb = MeshManager(port=port_b)
    _patch_config(ma)
    _patch_config(mb)

    app_a = FastAPI()
    app_a.include_router(build_mesh_router(ma))
    app_b = FastAPI()
    app_b.include_router(build_mesh_router(mb))

    sa = uvicorn.Server(uvicorn.Config(app_a, host="127.0.0.1", port=port_a, log_level="error"))
    sb = uvicorn.Server(uvicorn.Config(app_b, host="127.0.0.1", port=port_b, log_level="error"))
    ta = threading.Thread(target=sa.run, daemon=True)
    tb = threading.Thread(target=sb.run, daemon=True)
    try:
        ta.start()
        tb.start()
        if not _wait_listen("127.0.0.1", port_a) or not _wait_listen("127.0.0.1", port_b):
            _fail(failures, "live_handshake", "servers did not accept connections")
            return

        a_fp = public_key_b64_to_fingerprint(ma.identity.public_key_b64)
        mb.peers.remove_peer(a_fp)

        first = ma.connect_to_peer(address="127.0.0.1", port=port_b)
        if first.get("ok") is not False:
            _fail(failures, "live_handshake", f"first connect should be pending/fail ok, got {first}")
            return
        st = str(first.get("status") or "")
        if st != "pending":
            _fail(failures, "live_handshake", f"expected pending, got {first}")
            return

        pr = requests.get(f"http://127.0.0.1:{port_b}/mesh/peers/pending", timeout=5)
        pr.raise_for_status()
        pending = pr.json().get("pending") or []
        if not pending:
            _fail(failures, "live_handshake", "no pending requests on B")
            return
        rid = str(pending[0]["request_id"])
        rr = requests.post(
            f"http://127.0.0.1:{port_b}/mesh/peers/pending/{rid}",
            json={"action": "allow"},
            timeout=5,
        )
        rr.raise_for_status()
        if not rr.json().get("ok"):
            _fail(failures, "live_handshake", f"resolve failed: {rr.json()}")
            return

        second = ma.connect_to_peer(address="127.0.0.1", port=port_b)
        if not second.get("ok"):
            _fail(failures, "live_handshake", f"second connect failed: {second}")
            return
        if not ma.list_sessions():
            _fail(failures, "live_handshake", "A has no sessions after connect")
            return

        # One-way encrypted ping (peer handles ping without LLM).
        sess = ma.list_sessions()[0]
        sid = str(sess["session_id"])
        msg = MeshMessage(
            type="ping",
            conversation_id=f"ping_{uuid.uuid4().hex[:8]}",
            conversation_type="social",
            message_number=1,
            payload={"text": "p"},
            ttl=60,
        )
        send_out = ma.send_message(session_id=sid, message=msg, source="smoke")
        if not send_out.get("ok"):
            _fail(failures, "live_handshake", f"send_message failed: {send_out}")
            return

        _ok("live_handshake_connect", "handshake + ping delivered")
    except Exception as exc:
        _fail(failures, "live_handshake", str(exc))
    finally:
        try:
            sa.should_exit = True
            sb.should_exit = True
        except Exception:
            pass
        time.sleep(0.1)
        ma.stop()
        mb.stop()


def test_message_size_reject(failures: list[str]) -> None:
    """Theme 2: oversized cleartext message rejected before encrypt."""
    m = MeshManager(port=9205)
    _patch_config(m, message_max_bytes=8192)
    try:
        import os as _os

        session = __import__("mesh.session", fromlist=["MeshSession"]).MeshSession(
            session_id="s1",
            peer_fingerprint="fp",
            peer_hostname="h",
            peer_address="127.0.0.1",
            peer_port=9999,
            session_key=_os.urandom(32),
        )
        m.sessions["s1"] = session
        big_text = "Z" * 200_000
        msg = MeshMessage(
            type="prompt",
            conversation_id="big",
            conversation_type="social",
            message_number=1,
            payload={"text": big_text},
            ttl=60,
        )
        out = m.send_message(session_id="s1", message=msg)
        if out.get("ok"):
            _fail(failures, "message_size", "expected failure for huge message body")
            return
        _ok("message_size_reject", str(out.get("error", ""))[:80])
    finally:
        m.stop()


def main() -> int:
    failures: list[str] = []
    print("=== mesh_stack_smoke (local) ===\n")
    test_pending_per_peer_cap(failures)
    test_concurrent_register(failures)
    test_large_task_result_truncation(failures)
    test_session_cancel_pending(failures)
    test_message_size_reject(failures)
    test_live_handshake_connect(failures)

    print("\n=== Notes (theme 4 / real LAN) ===")
    print(
        "- Conflict resolution: mesh does not merge SESSION-STATE.md or MEMORY across machines; "
        "each peer has its own workspace. True shared-write tests need an explicit shared artifact "
        "(e.g. git repo, API) outside core mesh."
    )
    print(
        "- Concurrency storm via contact_peer: default pending_result_max_per_peer is 5 — "
        "sixth simultaneous wait-for-result returns 'too many outstanding mesh tasks for this peer'."
    )
    print(
        "- Full task_delegate round-trip requires a live LLM on the peer; not exercised here.\n"
    )

    if failures:
        print("=== SUMMARY: FAILED ===")
        for f in failures:
            print(" ", f)
        return 1
    print("=== SUMMARY: all checks passed ===")
    return 0


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