"""End-to-end test: Auth → WS Gateway → Chat Service round-trip.

Usage (from repo root, with venv activated):
    python test_e2e_flow.py

Requires: Auth (8001), WS Gateway (8002), Chat (8003) all running.
"""

import asyncio
import json
import uuid

import httpx
import websockets

AUTH_URL = "http://127.0.0.1:8001/api/v1/auth"
WS_URL = "ws://127.0.0.1:8002/api/v1/ws/device"

# ── 1. Authenticate ─────────────────────────────────────────────────


async def get_token() -> str:
    async with httpx.AsyncClient() as client:
        # Try login first, register if user doesn't exist
        resp = await client.post(
            f"{AUTH_URL}/login",
            json={"email": "e2e@test.com", "password": "Test1234!"},
        )
        if resp.status_code == 200:
            print("[1/4] Logged in as e2e@test.com")
            return resp.json()["access_token"]

        resp = await client.post(
            f"{AUTH_URL}/register",
            json={
                "email": "e2e@test.com",
                "password": "Test1234!",
                "name": "E2E",
                "surname": "Test",
            },
        )
        resp.raise_for_status()
        print("[1/4] Registered + logged in as e2e@test.com")
        return resp.json()["access_token"]


# ── 2. WebSocket flow ───────────────────────────────────────────────


async def run_e2e():
    token = await get_token()

    uri = f"{WS_URL}?token={token}"
    async with websockets.connect(uri) as ws:
        # Send device_hello
        await ws.send(json.dumps({
            "type": "device_hello",
            "device_id": str(uuid.uuid4()),
            "agent_ids": ["task", "note", "project", "timeline"],
        }))
        print("[2/4] Device registered with WS Gateway")

        # Send a home_request (simple greeting — unlikely to need tools)
        await ws.send(json.dumps({
            "type": "home_request",
            "message": "Hello! How are you doing today?",
            "context": {},
        }))
        print("[3/4] Sent home_request → waiting for Chat Service response...")

        # Listen for response frames (text_chunk, tool_call, final)
        full_response = []
        try:
            while True:
                raw = await asyncio.wait_for(ws.recv(), timeout=60)
                frame = json.loads(raw)
                ftype = frame.get("type")

                if ftype == "text_chunk":
                    chunk = frame.get("chunk", frame.get("text", ""))
                    full_response.append(chunk)
                    print(f"  ← text_chunk: {chunk[:80]}")

                elif ftype == "tool_call":
                    # Respond with a mock tool_result so the agent doesn't hang
                    call_id = frame.get("id")
                    action = frame.get("action")
                    table = frame.get("table", "")
                    print(f"  ← tool_call: {action} {table} (id={call_id})")

                    mock_result = {"rows": [], "row": None}
                    await ws.send(json.dumps({
                        "type": "tool_result",
                        "id": call_id,
                        **mock_result,
                    }))
                    print(f"  → tool_result (mock) for {call_id}")

                elif ftype == "final":
                    text = frame.get("text", "")
                    if text:
                        full_response.append(text)
                    print(f"  ← final")
                    break

                elif ftype == "ping":
                    # Ignore heartbeats
                    continue

                else:
                    print(f"  ← {ftype}: {json.dumps(frame)[:120]}")

        except asyncio.TimeoutError:
            print("  ⚠ Timed out waiting for response (60s)")

        print()
        if full_response:
            print(f"[4/4] Full response: {''.join(full_response)}")
        else:
            print("[4/4] No text response received (check Chat Service logs)")


if __name__ == "__main__":
    asyncio.run(run_e2e())