from typing import Dict, Any
import socket

import msgpack

import tcp_client_image

SERVER_ADDRESS = ('172.16.0.249', 8888)

def send_request(data:Dict[str, Any]):
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client_socket.connect(SERVER_ADDRESS)
    
    try:    
        message:bytes = msgpack.packb(data)

        client_socket.sendall(message)

        rx = client_socket.recv(1024)
        if data:
            print(f"Received message: {rx.decode('utf-8').strip()}")
        else:
            print("Connection closed")
            return
    except ConnectionResetError:
        print("Connection reset")
    finally:
        client_socket.close()
        print("Connection closed")

def image():
    image = b""
    try:
        choice = int(input("image > "))
        if choice == 1:
            image = tcp_client_image.image1
        elif choice == 2:
            image = tcp_client_image.image2
        else:
            raise ValueError
    except:
        print("invalid choice")
        return
    
    chunk_size = 2048
    for i in range(0, len(image), chunk_size):
        data = {
            "command": "update_image",
            "index": i,
            "data": bytes(image[i:i+chunk_size])
        }
        send_request(data)
    
    send_request({"command":"push_image"})

def main():
    while True:
        print("Command:")
        print("(1) update + push image")
        print("(2) ping")
        print("(3) invalid request")
        print("(4) exit")

        choice:int = int(input("> "))
        data = {}
        if choice == 1:
            image()
        elif choice == 2:
            send_request({"command":"ping"})
        elif choice == 3:
            send_request({"command":"invalid_command"})
        else:
            break

if __name__ == "__main__":
    main()