import socket
import struct
from time import sleep

def create_insert_message(timestamp: int, price: int) -> bytes:
    """Create a 9-byte insert message."""
    return struct.pack('!cii', b'I', timestamp, price)

def create_query_message(mintime: int, maxtime: int) -> bytes:
    """Create a 9-byte query message."""
    return struct.pack('!cii', b'Q', mintime, maxtime)

def receive_response(sock: socket.socket) -> int:
    """Receive and decode a 4-byte response."""
    data = sock.recv(4)
    return struct.unpack('!i', data)[0]

def run_test():
    # Connect to the server
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.connect(('localhost', 40000))  # Using the same port as your other servers
    
    # Test case 1: Insert some values
    test_data = [
        (12345, 101),  # (timestamp, price)
        (12346, 102),
        (12347, 100),
        (40960, 5),
    ]
    
    print("Inserting test data...")
    for timestamp, price in test_data:
        message = create_insert_message(timestamp, price)
        sock.send(message)
        print(f"Inserted: timestamp={timestamp}, price={price}")
        sleep(0.1)  # Small delay to avoid overwhelming the server
    
    # Test case 2: Query for mean price
    mintime = 12288
    maxtime = 16384
    print(f"\nQuerying mean price between {mintime} and {maxtime}...")
    
    query = create_query_message(mintime, maxtime)
    sock.send(query)
    
    # Get response
    mean_price = receive_response(sock)
    print(f"Received mean price: {mean_price}")
    
    # Test case 3: Query with no data in range
    print("\nTesting query with no data in range...")
    query = create_query_message(1, 1000)
    sock.send(query)
    mean_price = receive_response(sock)
    print(f"Received mean price (should be 0): {mean_price}")
    
    # Clean up
    sock.close()

if __name__ == "__main__":
    try:
        run_test()
    except ConnectionRefusedError:
        print("Could not connect to server. Make sure it's running on port 40000.")
    except Exception as e:
        print(f"An error occurred: {e}")