#!/usr/bin/env python3
"""
End-to-end LAN inference test for meshnet distributed inference.

Sends 3 chat-completion requests to a meshnet node, validates OpenAI-format
responses, and prints token counts + latency per request.

Usage:
    python scripts/test_lan_inference.py \\
        --tracker http://192.168.1.10:8080 \\
        --gateway http://192.168.1.10:8001

Exit 0 on success, 1 on any failure.
"""
from __future__ import annotations

import argparse
import json
import sys
import time
import urllib.error
import urllib.parse
import urllib.request


PROMPTS = [
    {"role": "user", "content": "What is 7 × 8? Answer in one word."},
    {"role": "user", "content": "Name the capital of France in one word."},
    {"role": "user", "content": "Complete the sequence: 1, 1, 2, 3, 5, ___. Answer in one word."},
]

MODEL = "microsoft/Phi-3-medium-128k-instruct"


def _get(url: str, timeout: float = 10.0) -> dict:
    with urllib.request.urlopen(url, timeout=timeout) as r:
        return json.loads(r.read())


def _post(url: str, payload: dict, timeout: float = 60.0) -> dict:
    data = json.dumps(payload).encode()
    req = urllib.request.Request(
        url, data=data, headers={"Content-Type": "application/json"}, method="POST"
    )
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())


def discover_gateway(tracker_url: str) -> str:
    """Return the first tracker-mode node endpoint for MODEL."""
    nodes = _get(f"{tracker_url}/v1/nodes", timeout=5.0)
    if isinstance(nodes, dict):
        nodes = list(nodes.values())
    tracker_nodes = [
        n for n in nodes
        if n.get("tracker_mode") and (
            n.get("hf_repo") == MODEL or n.get("model") == MODEL.split("/")[-1]
        )
    ]
    if not tracker_nodes:
        raise RuntimeError(
            f"No tracker-mode nodes found for {MODEL!r}. "
            "Is the first-shard node running and registered?"
        )
    endpoint: str = tracker_nodes[0]["endpoint"]
    return endpoint.rstrip("/")


def check_route(tracker_url: str, gateway_url: str) -> list[str]:
    """Return the full inference route for MODEL."""
    url = f"{tracker_url}/v1/route?model={urllib.parse.quote(MODEL)}"
    try:
        resp = _get(url, timeout=5.0)
        return resp.get("route", [])
    except Exception as exc:
        print(f"  Warning: could not fetch route: {exc}", file=sys.stderr)
        return [gateway_url]


def run_inference(gateway_url: str, messages: list[dict]) -> tuple[str, int, float]:
    """Send one chat-completion request. Returns (content, tokens, elapsed_s)."""
    t0 = time.monotonic()
    resp = _post(
        f"{gateway_url}/v1/chat/completions",
        {"model": MODEL, "messages": messages, "stream": False},
        timeout=120.0,
    )
    elapsed = time.monotonic() - t0

    choices = resp.get("choices")
    if not choices:
        raise ValueError(f"No choices in response: {resp}")
    content: str = choices[0].get("message", {}).get("content", "")
    if not isinstance(content, str):
        raise TypeError(f"Expected string content, got {type(content)}: {content}")

    usage = resp.get("usage", {})
    tokens: int = usage.get("completion_tokens", len(content.split()))

    return content, tokens, elapsed


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--tracker", required=True, help="Tracker URL, e.g. http://192.168.1.10:8080")
    p.add_argument(
        "--gateway",
        default=None,
        help="Inference entry point URL. Auto-discovered from tracker if omitted.",
    )
    args = p.parse_args(argv)

    tracker_url = args.tracker.rstrip("/")

    print(f"Tracker: {tracker_url}")

    # Resolve gateway
    gateway_url = args.gateway.rstrip("/") if args.gateway else None
    if gateway_url is None:
        try:
            gateway_url = discover_gateway(tracker_url)
            print(f"Gateway (auto-discovered): {gateway_url}")
        except Exception as exc:
            print(f"ERROR: {exc}", file=sys.stderr)
            return 1
    else:
        print(f"Gateway: {gateway_url}")

    # Show route
    route = check_route(tracker_url, gateway_url)
    print(f"Route: {route}")
    if len(route) < 2:
        print("  Warning: only one node in route — is the second-shard node registered?")
    print()

    failures = 0
    for i, msg in enumerate(PROMPTS, start=1):
        print(f"[{i}] Q: {msg['content']}")
        try:
            content, tokens, elapsed = run_inference(gateway_url, [msg])
            tps = tokens / elapsed if elapsed > 0 else 0.0
            print(f"     A: {content.strip()}")
            print(f"     {tokens} tokens  {elapsed:.2f}s  {tps:.1f} t/s")
        except urllib.error.HTTPError as exc:
            body = exc.read().decode(errors="replace")
            print(f"     ERROR {exc.code}: {body}", file=sys.stderr)
            failures += 1
        except Exception as exc:
            print(f"     ERROR: {exc}", file=sys.stderr)
            failures += 1
        print()

    if failures == 0:
        print(f"All {len(PROMPTS)} requests completed successfully.")
        print("Exit code: 0")
        return 0
    else:
        print(f"{failures}/{len(PROMPTS)} requests failed.", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())