import asyncio
import json
import websockets
import pytest
from pydantic import BaseModel

async def test_websocket_sync():
    uri = "ws://localhost:8000/ws"
    async with websockets.connect(uri) as websocket:
        # 1. Send hello
        hello = {
            "type": "hello",
            "seqs": {}
        }
        await websocket.send(json.dumps(hello))
        
        # 2. Receive snapshots
        # Expecting TraderState and StrategyState
        responses = []
        for _ in range(2):
            resp = await websocket.recv()
            responses.append(json.loads(resp))
            
        assert any(r["store"] == "TraderState" for r in responses)
        assert any(r["store"] == "StrategyState" for r in responses)
        
        # 3. Send a patch for TraderState
        trader_resp = next(r for r in responses if r["store"] == "TraderState")
        current_seq = trader_resp["seq"]
        
        patch_msg = {
            "type": "patch",
            "store": "TraderState",
            "seq": current_seq,
            "patch": [{"op": "replace", "path": "/status", "value": "busy"}]
        }
        await websocket.send(json.dumps(patch_msg))
        
        # 4. Receive confirmation patch
        confirm_resp = await websocket.recv()
        confirm_json = json.loads(confirm_resp)
        assert confirm_json["type"] == "patch"
        assert confirm_json["store"] == "TraderState"
        assert confirm_json["seq"] == current_seq + 1
        assert confirm_json["patch"][0]["value"] == "busy"

if __name__ == "__main__":
    # This script requires the server to be running:
    # PYTHONPATH=backend/src python3 backend/src/main.py
    try:
        asyncio.run(test_websocket_sync())
        print("Test passed!")
    except Exception as e:
        print(f"Test failed: {e}")