import os
import struct
from socket import *
import sys

def mainMenu():
    print("1. Nastavit PC ako klienta")
    print("2. Nastavit PC ako server")
    print("3. Vypnut")
    selection = int(input("Vyber:"))
    if selection==1:
        start_client()
    elif selection==2:
        start_server()
    elif selection==3:
        exit
    else:
        print("Vyber cislo z menu!")
        mainMenu()

def start_client():
    content = None
    def client_menu1():
        host = input("Zadaj IP adresu kam chces posielat v tvare xxx.xxx.xxx.xxx\n")
        return host

    # file_name = print("Zadaj nazov suboru aj s cestou\n")
    # with open('file_name', 'rb') as f:
    #     content = f.read()
    host = client_menu1()
    # host = "127.0.0.1"  # set to IP address of target computer
    print(host)
    port = 10051
    addr = (host, port)
    UDPSock = socket(AF_INET, SOCK_DGRAM)

    def header_connection_msg():
        type = 0 #0 connection, 1 - info, 2 - data_head
        size_fragments = 1500#user_edit
        num_fragments = 1 #sizeofMSG/sizeoffragments
        type_of_msg = 0 #0 - text, 1 - file
        #checksum = CRC
        header = struct.pack("IHHI", type, size_fragments, num_fragments, type_of_msg)
        return header

    datap = 0

    if UDPSock.sendto(header_connection_msg(), addr):
        (datap, addr) = UDPSock.recvfrom(1024)

    if datap == "0":
        print("Nic")
    else:
        print("OK")

    while True:
         data = input("Zadaj spravu")
         UDPSock.sendto(data.encode(), addr)
         (datap, addr) = UDPSock.recvfrom(1024)
         print(datap.decode())
         if data == "exit":
             break

    UDPSock.close()
    os._exit(0)


def start_server():
    host = "127.0.0.1"
    port = 10051
    buf = 1024

    address = (host, port)
    UDPSock = socket(AF_INET, SOCK_DGRAM)
    (header, address) = UDPSock.recvfrom(1024)
    type, size_fragments, num_fragments, type_of_msg = struct.unpack("cHHc", header)
    p = 0
    if type == 0:
        p = "0"
    UDPSock.sendto(p.encode(), address)
    UDPSock.bind(address)
    print('Server caka')

    while True:
        (data, address) = UDPSock.recvfrom(5)
        data1 = "Server prijal spravu"
        UDPSock.sendto(data1.encode(), address)
        new_data = data.decode()
        print(f"Received message: {sys.getsizeof(data)}")
        print(f"Received message: {data}")
        if data == b"exit":
            break

    UDPSock.close()
    os._exit(0)

mainMenu()