"""E2E check: two real TorchNodeServers over HTTP, user's live topology.

Head owns layers 0-20, tail owns 10-23 (overlapping shards, clamped to @21),
exactly like the reported 2-node run. Compares cached vs stateless tps and
verifies identical output text.
"""
import json
import sys
import time
import urllib.request

sys.path.insert(0, "packages/node")

from meshnet_node.model_backend import TorchModelShard
from meshnet_node.torch_server import TorchNodeServer

MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
MAX_TOKENS = 60

print("loading shards...", flush=True)
head_backend = TorchModelShard(MODEL, 0, 20)
tail_backend = TorchModelShard(MODEL, 10, 23)

head = TorchNodeServer(backend=head_backend, tracker_mode=True)
tail = TorchNodeServer(backend=tail_backend)
head_port = head.start()
tail_port = tail.start()
print(f"head :{head_port} layers 0-20 | tail :{tail_port} layers 10-23", flush=True)


def chat(label):
    body = json.dumps({
        "model": MODEL,
        "messages": [{"role": "user", "content": "Count from 1 to 10, then say done."}],
        "max_tokens": MAX_TOKENS,
        "stream": False,
    }).encode()
    req = urllib.request.Request(
        f"http://127.0.0.1:{head_port}/v1/chat/completions",
        data=body,
        headers={
            "Content-Type": "application/json",
            # tracker-injected route: tail continues at layer 21 (clamped)
            "X-Meshnet-Route": json.dumps(
                [{"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 21}]
            ),
        },
        method="POST",
    )
    started = time.monotonic()
    with urllib.request.urlopen(req, timeout=600) as r:
        resp = json.loads(r.read())
    elapsed = time.monotonic() - started
    text = resp["choices"][0]["message"]["content"]
    tokens = resp.get("usage", {}).get("completion_tokens", 0)
    print(f"\n[{label}] tokens={tokens} elapsed={elapsed:.2f}s "
          f"tps={tokens / max(elapsed, 1e-6):.2f}", flush=True)
    print(f"[{label}] text={text!r}", flush=True)
    return text, tokens, elapsed


try:
    cached_text, ct, ce = chat("cached-warmup")  # first run pays model warmup
    cached_text, ct, ce = chat("cached")
    forwards_cached = tail.forward_chunk_count

    head_backend.supports_kv_cache = False  # force legacy stateless path
    stateless_text, st, se = chat("stateless")
    forwards_stateless = tail.forward_chunk_count - forwards_cached

    print(f"\ntail /forward calls: cached-run total={forwards_cached} "
          f"stateless-run={forwards_stateless}", flush=True)
    print(f"speedup: {(ct / ce) / max(st / se, 1e-9):.2f}x", flush=True)
    print(f"outputs identical: {cached_text == stateless_text}", flush=True)
finally:
    head.stop()
    tail.stop()