#!/usr/bin/env python3
import socket
import struct
import msgpack
import sys
import os
import time

# Define test data: 5 requests with different pinyin and context
test_requests = [
    {"pinyin": "shanghai", "context": ""},
    {"pinyin": "beijing", "context": ""},
    {"pinyin": "nihao", "context": ""},
    {"pinyin": "xiexie", "context": ""},
    {"pinyin": "f", "context": "返回汉字的完整拼音和剩余部"}
]

# Determine socket address based on operating system
def get_socket_address():
    if sys.platform == "win32":
        # Windows: TCP socket
        return ("127.0.0.1", 23333), "tcp"
    else:
        # Unix-like systems: Unix Domain Socket
        return "/tmp/su-ime.sock", "unix"

# Connect to server
def connect_to_server(address, socket_type):
    if socket_type == "unix":
        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        sock.connect(address)
    else:  # tcp
        host, port = address
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.connect((host, port))
    return sock

# Send request and receive response
def send_request(sock, request):
    # Serialize request to MessagePack
    payload = msgpack.packb(request)
    payload_len = len(payload)
    # Create header: 4-byte big-endian unsigned integer
    header = struct.pack(">I", payload_len)
    # Send header and payload
    sock.sendall(header + payload)
    
    # Receive response
    # Read header
    len_buf = sock.recv(4)
    if len(len_buf) != 4:
        raise ValueError("Failed to read response header")
    response_len = struct.unpack(">I", len_buf)[0]
    
    # Read payload
    response_payload = b""
    while len(response_payload) < response_len:
        chunk = sock.recv(response_len - len(response_payload))
        if not chunk:
            break
        response_payload += chunk
    
    # Deserialize response
    response = msgpack.unpackb(response_payload, raw=False)
    return response

# Main function
def main():
    address, socket_type = get_socket_address()
    
    try:
        sock = connect_to_server(address, socket_type)
        print(f"Connected to server via {socket_type} at {address}")
        
        for i, req in enumerate(test_requests, 1):
            print(f"\nSending request {i}: {req}")
            try:
                start = time.time()
                response = send_request(sock, req)
                print(f'cost time: {time.time() - start}')
                print(f"Response {i}: {response[0][0:10]}")
            except Exception as e:
                print(f"Error sending request {i}: {e}")
        
        sock.close()
        print("\nAll requests sent.")
    except Exception as e:
        print(f"Connection error: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()