#!/usr/bin/env python3
"""End-to-end message delivery verification for GoChat.

Tests:
1. Agent sends message via REST API → WebSocket receives message.created event
2. Measures delivery latency (P50/P95/P99) and success rate over N messages
3. Verifies event payload structure (event type, account_id, data.message)

Usage:
    python3 verify_msg_e2e.py [--host 127.0.0.1] [--port 3000] [--count 50]
"""

import argparse
import asyncio
import json
import statistics
import sys
import time
import urllib.request

try:
    import websockets
except ImportError:
    print("ERROR: pip install websockets", file=sys.stderr)
    sys.exit(1)


def login(host, port, email, password):
    """Login and return Chatwoot auth headers (JWT access-token + client + uid)."""
    url = f"http://{host}:{port}/auth/sign_in"
    body = json.dumps({"email": email, "password": password}).encode()
    req = urllib.request.Request(
        url, data=body, method="POST",
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=10) as resp:
        headers = {k.lower(): v for k, v in resp.headers.items()}
        body = json.loads(resp.read())
    return {
        "access-token": headers.get("access-token", ""),
        "client": headers.get("client", ""),
        "uid": headers.get("uid", ""),
        "token-type": headers.get("token-type", "Bearer"),
        "pubsub_token": body.get("data", {}).get("pubsub_token", ""),
        "account_id": body.get("data", {}).get("account_id", 1),
        "user": body.get("data", {}),
    }


def get_conversations(host, port, auth, account_id):
    """List conversations to find one for testing."""
    url = f"http://{host}:{port}/api/v1/accounts/{account_id}/conversations"
    req = urllib.request.Request(url, method="GET", headers=auth_headers(auth))
    with urllib.request.urlopen(req, timeout=10) as resp:
        return json.loads(resp.read())


def auth_headers(auth):
    return {
        "access-token": auth["access-token"],
        "client": auth["client"],
        "uid": auth["uid"],
        "token-type": auth["token-type"],
        "Content-Type": "application/json",
    }


def send_message(host, port, auth, account_id, conversation_id, content):
    """POST a message via REST API, return (send_time, response_body)."""
    url = (
        f"http://{host}:{port}/api/v1/accounts/{account_id}"
        f"/conversations/{conversation_id}/messages"
    )
    body = json.dumps({
        "content": content,
        "content_type": "text",
        "message_type": "outgoing",
        "private": False,
    }).encode()
    req = urllib.request.Request(url, data=body, method="POST", headers=auth_headers(auth))
    send_time = time.monotonic()
    with urllib.request.urlopen(req, timeout=10) as resp:
        resp_body = json.loads(resp.read())
    return send_time, resp_body


async def ws_connect(host, port, auth, account_id):
    """Connect to WebSocket /cable and subscribe to AccountChannel.

    Uses ActionCable-compatible command format:
      {"command":"subscribe","identifier":"{\"channel\":\"AccountChannel\",\"account_id\":1}"}
    The client is auto-subscribed to the account room on connect (hub.go:180),
    but we send an explicit subscribe for confirmation.
    """
    ws_url = f"ws://{host}:{port}/cable?token={auth['access-token']}"
    ws = await websockets.connect(ws_url, max_size=2**20)
    # Subscribe to AccountChannel (ActionCable format: identifier is JSON string)
    identifier = json.dumps({"channel": "AccountChannel", "account_id": account_id})
    sub_cmd = {"command": "subscribe", "identifier": identifier}
    await ws.send(json.dumps(sub_cmd))
    # Wait for subscribe confirmation or welcome
    while True:
        raw = await asyncio.wait_for(ws.recv(), timeout=5)
        msg = json.loads(raw)
        msg_type = msg.get("type", "")
        if msg_type == "confirm_subscribe":
            return ws
        if msg.get("event") == "welcome":
            # Auto-subscribed to account room; keep waiting for confirm
            continue
        if msg_type == "reject_subscribe":
            raise RuntimeError(f"subscribe rejected: {msg.get('reason')}")
        # Unknown message — assume connected (auto-subscribe may skip confirm)
        return ws


async def verify_message_delivery(host, port, count):
    print(f"\n=== GoChat Message E2E Verification ===")
    print(f"Target: http://{host}:{port}")
    print(f"Messages: {count}")

    # Step 1: Login
    print("\n[1/4] Logging in as admin...")
    auth = login(host, port, "admin@gochat.local", "changeme")
    account_id = auth["account_id"]
    print(f"  OK: user={auth['user']['name']}, account_id={account_id}")

    # Step 2: Find a conversation
    print("[2/4] Finding conversation...")
    convos = get_conversations(host, port, auth, account_id)
    payload = convos.get("data", {}).get("payload", [])
    if not payload:
        print("  FAIL: no conversations found. Run `gochat seed` first.")
        return False
    conv = payload[0]
    conv_id = conv["id"]
    sender = conv.get("meta", {}).get("sender", {}).get("name", "?")
    print(f"  OK: conversation_id={conv_id} ({sender})")

    # Step 3: Connect WebSocket
    print("[3/4] Connecting WebSocket...")
    try:
        ws = await ws_connect(host, port, auth, account_id)
    except Exception as e:
        print(f"  FAIL: WS connect/subscribe failed: {e}")
        return False
    print(f"  OK: subscribed to AccountChannel")

    # Step 4: Send messages and measure latency
    print(f"[4/4] Sending {count} messages and measuring delivery latency...")
    latencies = []
    received = 0
    missed = 0

    for i in range(count):
        content = f"e2e-test-{i:04d}-{int(time.time()*1000)}"
        try:
            send_time, resp = send_message(host, port, auth, account_id, conv_id, content)
            msg_id = resp.get("id")
        except Exception as e:
            print(f"  [{i+1}] SEND FAIL: {e}")
            missed += 1
            continue

        # Wait for message.created event on WS
        try:
            while True:
                raw = await asyncio.wait_for(ws.recv(), timeout=5)
                msg = json.loads(raw)
                if msg.get("event") != "message.created":
                    continue
                # EventFrame uses 'payload' for data, WSMessage uses 'data'
                data = msg.get("payload") or msg.get("data", {})
                # message id is nested under data.message.id (dispatchMessageEvent
                # sets event.Data["message"] = message)
                msg_data = data.get("message", data)
                if msg_data.get("id") == msg_id:
                    recv_time = time.monotonic()
                    latency_ms = (recv_time - send_time) * 1000
                    latencies.append(latency_ms)
                    received += 1
                    break
                # Different message — might be echo of our own or other; keep waiting
        except asyncio.TimeoutError:
            print(f"  [{i+1}] TIMEOUT: no WS event for msg_id={msg_id}")
            missed += 1

        if (i + 1) % 10 == 0:
            print(f"  ...{i+1}/{count} sent, {received} received")

    await ws.close()

    # Results
    print(f"\n=== Results ===")
    print(f"Sent:      {count}")
    print(f"Received:  {received}")
    print(f"Missed:    {missed}")
    if latencies:
        latencies.sort()
        p50 = latencies[len(latencies) // 2]
        p95 = latencies[int(len(latencies) * 0.95)]
        p99 = latencies[min(int(len(latencies) * 0.99), len(latencies) - 1)]
        mean = statistics.mean(latencies)
        print(f"\nLatency (ms):")
        print(f"  Mean:  {mean:.1f}")
        print(f"  P50:   {p50:.1f}")
        print(f"  P95:   {p95:.1f}")
        print(f"  P99:   {p99:.1f}")
        print(f"  Min:   {latencies[0]:.1f}")
        print(f"  Max:   {latencies[-1]:.1f}")
        print(f"\nSuccess rate: {received/count*100:.1f}%")
    return received == count


def main():
    parser = argparse.ArgumentParser(description="GoChat message E2E verification")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=3000)
    parser.add_argument("--count", type=int, default=50, help="number of messages to send")
    args = parser.parse_args()

    ok = asyncio.run(verify_message_delivery(args.host, args.port, args.count))
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()