#!/usr/bin/env python3
"""Quick test to verify state recovery handshake protocol implementation."""

import asyncio

from src.services.websocket_service import WebSocketService


async def test_state_recovery():
    """Test the state recovery components."""

    # Create service
    service = WebSocketService()

    print("Testing State Recovery Components")
    print("=" * 50)

    # Test 1: Message buffer initialization
    print("\n1. Message buffer initialization:")
    print(f"   - Buffer type: {type(service.message_buffer)}")
    print(f"   - Buffer maxlen: {service.message_buffer.maxlen}")
    print(f"   - Current sequence: {service.current_sequence}")
    print("   ✓ Buffer initialized correctly")

    # Test 2: Add sequence to messages
    print("\n2. Adding sequence to messages:")
    msg1 = {"type": "console_log", "data": "test1"}
    msg2 = {"type": "console_log", "data": "test2"}
    msg3 = {"type": "console_log", "data": "test3"}

    sequenced_msg1 = service._add_sequence(msg1.copy())
    sequenced_msg2 = service._add_sequence(msg2.copy())
    sequenced_msg3 = service._add_sequence(msg3.copy())

    print(f"   - Message 1 sequence: {sequenced_msg1.get('sequence')}")
    print(f"   - Message 2 sequence: {sequenced_msg2.get('sequence')}")
    print(f"   - Message 3 sequence: {sequenced_msg3.get('sequence')}")
    print(f"   - Current sequence: {service.current_sequence}")
    print(f"   - Buffer size: {len(service.message_buffer)}")
    print("   ✓ Sequences added correctly")

    # Test 3: Get messages after sequence
    print("\n3. Getting messages after sequence:")
    replay_from_0 = service._get_messages_after_sequence(0)
    replay_from_1 = service._get_messages_after_sequence(1)
    replay_from_2 = service._get_messages_after_sequence(2)

    print(f"   - Messages after sequence 0: {len(replay_from_0)} messages")
    print(f"   - Messages after sequence 1: {len(replay_from_1)} messages")
    print(f"   - Messages after sequence 2: {len(replay_from_2)} messages")

    assert len(replay_from_0) == 3, "Should have 3 messages after sequence 0"
    assert len(replay_from_1) == 2, "Should have 2 messages after sequence 1"
    assert len(replay_from_2) == 1, "Should have 1 message after sequence 2"
    print("   ✓ Message replay works correctly")

    # Test 4: Connection init message structure
    print("\n4. Connection init message structure:")
    connection_init = {
        "type": "connection_init",
        "lastSequence": 1,
        "extensionVersion": "1.0.0",
        "capabilities": ["console", "dom"],
    }
    print(f"   - Message type: {connection_init.get('type')}")
    print(f"   - Last sequence: {connection_init.get('lastSequence')}")
    print(f"   - Extension version: {connection_init.get('extensionVersion')}")
    print(f"   - Capabilities: {connection_init.get('capabilities')}")
    print("   ✓ Message structure is correct")

    # Test 5: Connection ack structure
    print("\n5. Connection ack response structure:")
    try:
        from src._version import __version__

        server_version = __version__
    except ImportError:
        server_version = "2.1.1"

    ack_message = {
        "type": "connection_ack",
        "serverVersion": server_version,
        "currentSequence": service.current_sequence,
        "replay": replay_from_1[:100],
    }

    print(f"   - Response type: {ack_message.get('type')}")
    print(f"   - Server version: {ack_message.get('serverVersion')}")
    print(f"   - Current sequence: {ack_message.get('currentSequence')}")
    print(f"   - Replay messages: {len(ack_message.get('replay', []))} messages")
    print("   ✓ Connection ack structure is correct")

    # Test 6: Buffer limit (maxlen=1000)
    print("\n6. Testing buffer limit (1000 messages):")
    for i in range(1005):
        service._add_sequence({"type": "test", "data": f"msg_{i}"})

    print("   - Messages added: 1005")
    print(f"   - Buffer size: {len(service.message_buffer)}")
    print(f"   - Sequence incremented to: {service.current_sequence}")
    assert len(service.message_buffer) == 1000, "Buffer should be capped at 1000"
    print("   ✓ Buffer limit works correctly")

    print("\n" + "=" * 50)
    print("All tests passed! ✓")
    print("\nState Recovery Handshake Protocol is ready to use.")
    print("\nNext steps:")
    print("  1. Start the MCP server")
    print("  2. Reload the extension")
    print("  3. Check server logs for 'Connection init' messages")
    print("  4. Verify 'connection_ack' responses with replay data")


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