import asyncio
import json
from websockets import connect
import httpx

async def test():
    print("=== Test Auto-Connect ===", flush=True)

    async with httpx.AsyncClient() as client:
        r = await client.post('http://localhost:3030/auth/login',
                              json={'username': 'admin', 'password': 'admin'})
        token = r.json()['token']

    async with connect('ws://localhost:3030/ws/chat') as ws:
        await ws.send(json.dumps({'token': token}))
        init = json.loads(await ws.recv())
        print(f"1. Init: {init.get('type')}, sessions: {len(init.get('sessions', []))}", flush=True)

        # Create session
        print("2. Creating session...", flush=True)
        await ws.send(json.dumps({'type': 'create_session', 'name': 'mi_sesion'}))

        # Should receive session_created AND session_connected
        for i in range(3):
            msg = json.loads(await asyncio.wait_for(ws.recv(), timeout=5))
            print(f"   Received: {msg.get('type')} - {msg}", flush=True)
            if msg.get('type') == 'session_connected':
                print("3. AUTO-CONNECT OK!", flush=True)

                # Now send a message
                print("4. Sending message...", flush=True)
                await ws.send(json.dumps({'type': 'message', 'content': 'di hola'}))

                for j in range(30):
                    try:
                        resp = json.loads(await asyncio.wait_for(ws.recv(), timeout=2))
                        print(f"   Response: {resp}", flush=True)
                        if resp.get('type') == 'done':
                            print("\n=== SUCCESS ===", flush=True)
                            return
                    except asyncio.TimeoutError:
                        pass
                return

asyncio.run(test())