import json
from multiprocessing.connection import Client
import pickle
import random
import socket
import sys
import time

from message import *  # Can change this after testing done
from timers import ClientTimer, setTimer


class LockClient():
    def __init__(self) -> None:
        with open("config.json", encoding='utf-8') as f:
            json_dict = json.loads(f.read())
        server_addresses = json_dict["servers"]
        self.server_addresses = [(server_addresses[i], int(server_addresses[i + 1])) for i in range(0, len(server_addresses), 2)]

        # create a UDP socket
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock.bind(('', 0))
        self.client_addr = self.sock.getsockname()
        print(f"Starting client {self.client_addr}", file=sys.stdout)
        random.seed(0)
        # For client timer to know whether to resend message
        self.req_num = 0
        self.num_requests = json_dict["num_request_per_client"]
        self.time_limit = json_dict["time_limit"]
        self.time_end = 0

    def executeLatency(self):
        lock_cmd = ["lock", "unlock"]
        print("Start sending requests..")
        latency_list = list()
        try:
            for i in range(0, self.num_requests):
                lock = random.choice(lock_cmd)
                val = random.randrange(1000)

                paxos_req = PaxosRequest(self.client_addr, LockCommand(lock, val, self.req_num))
                data = pickle.dumps(paxos_req)

                server_address = random.choice(self.server_addresses)
                print(f"Sending command {paxos_req} to {server_address}", file=sys.stdout)
                start_time = time.time()
                self.sock.sendto(data, server_address)

                setTimer(ClientTimer(paxos_req, self.req_num), ClientTimer.CLIENT_RETRY_MILLIS, self.onClientTimer)

                # receive data back from the server
                res = pickle.loads(self.sock.recv(1024))
                while res.cmd.req_num != self.req_num:
                    res = pickle.loads(self.sock.recv(1024))

                end_time = time.time()
                latency = end_time - start_time
                latency_list.append(latency)
                self.req_num += 1

        finally:
            # Close socket
            self.sock.close()

            print(f"\nFinished running.. Sent {self.req_num} out of {self.num_requests} total requests")
            print(f"\nResults")
            print(f"Average Latency: {sum(latency_list)/len(latency_list)}")

    def executeThroughput(self):
        lock_cmd = ["lock", "unlock"]
        print("Start sending requests..")
        self.req_num = 0
        self.time_end = time.time() + self.time_limit
        try:
            while time.time() < self.time_end:
                lock = random.choice(lock_cmd)
                val = random.randrange(1000)

                paxos_req = PaxosRequest(self.client_addr, LockCommand(lock, val, self.req_num))
                data = pickle.dumps(paxos_req)

                server_address = random.choice(self.server_addresses)
                print(f"Sending command {paxos_req} to {server_address}", file=sys.stdout)
                self.sock.sendto(data, server_address)

                setTimer(ClientTimer(paxos_req, self.req_num), ClientTimer.CLIENT_RETRY_MILLIS, self.onClientTimerV2)

                # receive data back from the server
                res = pickle.loads(self.sock.recv(1024))
                while res.cmd.req_num != self.req_num:
                    if time.time() >= self.time_end:
                        break
                    res = pickle.loads(self.sock.recv(1024))

                self.req_num += 1

        finally:
            # Close socket
            self.sock.close()

            print(f"\nFinished running.. Sent {self.req_num} out of {self.num_requests} total requests")
            print(f"\nResults")
            print(f"Throughput: {self.req_num / self.time_limit}")

    def onClientTimer(self, client_timer: ClientTimer):
        # print(f"{client_timer}: Callback", file=sys.stdout)
        if self.req_num <= client_timer.req_num:
            server_address = random.choice(self.server_addresses)
            print(f"Resending command {client_timer.paxos_req} to {server_address}", file=sys.stdout)
            data = pickle.dumps(client_timer.paxos_req)
            self.sock.sendto(data, server_address)
            setTimer(client_timer, ClientTimer.CLIENT_RETRY_MILLIS, self.onClientTimer)

    def onClientTimerV2(self, client_timer: ClientTimer):
        # print(f"{client_timer}: Callback", file=sys.stdout)
        if self.req_num <= client_timer.req_num and time.time() < self.time_end:
            server_address = random.choice(self.server_addresses)
            print(f"Resending command {client_timer.paxos_req} to {server_address}", file=sys.stdout)
            data = pickle.dumps(client_timer.paxos_req)
            self.sock.sendto(data, server_address)
            setTimer(client_timer, ClientTimer.CLIENT_RETRY_MILLIS, self.onClientTimer)


if __name__ == "__main__":
    client = LockClient()
    client.executeLatency()
    #client.executeThroughput()