import asyncio
import json
from websockets import connect
import httpx

async def test():
    print("=== Test Connect Session ===")

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

    # Connect and create session
    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 sessions: {len(init.get('sessions', []))}")
        for s in init.get('sessions', []):
            print(f"   - {s['name']} ({s['session_id'][:8]}...)")

        # Create a new session
        await ws.send(json.dumps({'type': 'create_session', 'name': 'test_session'}))
        created = json.loads(await ws.recv())
        print(f"2. Created: {created}")
        sid = created.get('session_id')

        # Should auto-connect - check for session_connected
        connected = json.loads(await ws.recv())
        print(f"3. Connected msg: {connected}")

    print("\n--- Reconnecting ---\n")

    # Reconnect
    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"4. Init sessions: {len(init.get('sessions', []))}")
        for s in init.get('sessions', []):
            print(f"   - {s['name']} ({s['session_id'][:8]}...)")

        # Try to connect to the session we created
        print(f"\n5. Connecting to session {sid[:8]}...")
        await ws.send(json.dumps({'type': 'connect_session', 'session_id': sid}))
        result = json.loads(await ws.recv())
        print(f"6. Result: {result}")

        if result.get('type') == 'session_connected':
            print("\n=== SUCCESS - Connected to existing session ===")

            # Try sending a message
            await ws.send(json.dumps({'type': 'message', 'content': 'test'}))
            for i in range(30):
                try:
                    m = await asyncio.wait_for(ws.recv(), timeout=1.0)
                    print(f"   RECV: {json.loads(m)}")
                    if json.loads(m).get('type') == 'done':
                        break
                except asyncio.TimeoutError:
                    pass
        else:
            print(f"\n=== FAILED: {result} ===")

asyncio.run(test())