#!/usr/bin/env python3
"""Test script for Pixel Nano (Gemma 3n on-device model via OpenAI-compatible API).

Usage:
    python3 test-pixel-nano.py [--host HOST] [--port PORT]
"""

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


def api_get(base_url: str, path: str) -> dict:
    req = urllib.request.Request(f"{base_url}{path}", method="GET")
    with urllib.request.urlopen(req, timeout=10) as resp:
        return json.loads(resp.read())


def chat(base_url: str, prompt: str, max_tokens: int = 512) -> dict:
    body = json.dumps({
        "model": "pixel10",
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
    }).encode()
    req = urllib.request.Request(
        f"{base_url}/v1/chat/completions",
        data=body,
        method="POST",
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=60) as resp:
        return json.loads(resp.read())


def main():
    parser = argparse.ArgumentParser(description="Test Pixel Nano API")
    parser.add_argument("--host", default="192.168.1.65", help="Pixel IP (default: 192.168.1.65)")
    parser.add_argument("--port", type=int, default=8080, help="API port (default: 8080)")
    parser.add_argument("--prompt", default=None, help="Custom prompt")
    args = parser.parse_args()

    base_url = f"http://{args.host}:{args.port}"
    print(f"Pixel Nano test — {base_url}\n")

    # 1) Health
    print("=== Health ===")
    try:
        h = api_get(base_url, "/health")
        print(f"  Device:  {h.get('device')}")
        print(f"  Model:   {h.get('model')}")
        print(f"  Status:  {h.get('status')}")
        print(f"  Uptime:  {h.get('uptime_seconds')}s")
        print(f"  Served:  {h.get('requests_served')} requests")
    except Exception as e:
        print(f"  FAILED: {e}")
        print("  Is PixelNano app running on the phone?")
        sys.exit(1)

    if h.get("status") != "ready":
        print("\nModel not ready — aborting.")
        sys.exit(1)

    # 2) Models
    print("\n=== Models ===")
    try:
        m = api_get(base_url, "/v1/models")
        for model in m.get("data", []):
            print(f"  {model['id']}: ctx={model.get('context_length')}, max_out={model.get('max_output_tokens')}")
    except Exception as e:
        print(f"  FAILED: {e}")

    # 3) Inference
    prompts = [args.prompt] if args.prompt else [
        "What is 2 + 2?",
        "Write a Python one-liner that prints 'Hello World'.",
        "Explain gravity in one sentence.",
    ]

    print("\n=== Inference ===")
    for i, prompt in enumerate(prompts, 1):
        print(f"\n  [{i}] {prompt}")
        try:
            t0 = time.time()
            result = chat(base_url, prompt)
            wall_ms = (time.time() - t0) * 1000

            msg = result.get("choices", [{}])[0].get("message", {}).get("content", "???")
            usage = result.get("usage", {})
            print(f"      → {msg}")
            print(f"      Wall: {wall_ms:.0f}ms | Tokens: {usage.get('prompt_tokens', '?')}→{usage.get('completion_tokens', '?')}")
        except urllib.error.HTTPError as e:
            print(f"      ERROR {e.code}: {e.read().decode()}")
        except Exception as e:
            print(f"      FAILED: {e}")

    print("\nDone!")


if __name__ == "__main__":
    main()