
"""
\x04: end of message
\x01: close requested
\x02: end of request

"""


class Socket():
    """A class of functions to make making tcp servers eazier
    """
    def __init__(self):
        import threading
        
        self.exceptions=self._exceptions() # initialize exceptions
        self.mode="unknown" # set modes to unused
        self.status="preinit"
        self.closing=threading.Event()
        self.needsClosing=0
        self.clientList=[]

    class _exceptions():
        class AlreadyInitializedException(Exception):
            "The socket has already been initialized"
            pass


    # mode setting functions
    
    def server(self,port: int,ip: str=""):
        """Make this socket a server.
        mode=server

        Args:
            port (int): the port to listen on
            ip (str): the ip address to listen on. Use `""` for nothing

        Raises:
            self.exceptions.AlreadyInitializedException: This means your socket already has a mode set.
        """
        
        if not self.status=="preinit": # check if socket is already started
            raise self.exceptions.AlreadyInitializedException("Can't initialize a socket twice")
        from socket import socket
        self.mode="server" # set modes
        self.status="init"

        self.sock=socket() # init socket
        self.sock.bind((ip,port)) # tell socket where to listen
        self.sock.listen() # set socket to listen mode

    def client(self,ip: str,port: int):
        """Make this socket a client

        Args:
            ip (str): the ip of the server to connect to
            port (int): the port of the server to connect to

        Raises:
            self.exceptions.AlreadyInitializedException: This means your socket already has a mode set.
        """
        
        if not self.status=="preinit": # check if socket is already started
            raise self.exceptions.AlreadyInitializedException("Can't initialize a socket twice")
        from socket import socket
        self.mode="client" # set modes
        self.status="init"

        self.sock=socket() # init socket
        self.sock.connect((ip,port)) # connect to server

    def setsocket(self,sock):
        """Make this a modeless socket. Can be used for the sockets of clients on a server.
        Also usefull if you just want the protocol functions, not the server management.
        
        Args:
            sock (socket): the socket (from socket.socket()) to use.

        Raises:
            self.exceptions.AlreadyInitializedException: This means your socket already has a mode set.
        """
        
        if not self.status=="preinit": # check if socket is already started
            raise self.exceptions.AlreadyInitializedException("Can't initialize a socket twice")
        self.mode="serverClient"
        self.status="handeled"
        self.sock=sock

    # server specific functions
    
    def handleClients(self,fn: object|None =None):
        import threading
        self.needsClosing=self.needsClosing+1 # register new loop
        try:
            def wrapper(fn,client,addr): # wrapper function for clients
                # register client
                self.clientList.append({"sock": client, "addr": addr})
                self.needsClosing=self.needsClosing+1
                # run client functions
                try:
                    fn(client,addr)
                except Exception as e: # catch exception
                    import traceback
                    tb=traceback.format_exc() # get stacktrace
                    print(tb) # print stacktrace
                # unregister client
                self.needsClosing=self.needsClosing-1
                self.clientList.remove({"sock": client, "addr": addr})
                
            while not self.closing.is_set(): # while not in closing state
                client,addr=self.sock.accept() # accept new client
                clientth=threading.Thread(target=wrapper,args=[fn,client,addr])
                clientth.start() # start client function
                
        except Exception as e: # catch exception
            import traceback
            tb=traceback.format_exc() # get stacktrace
            print(tb) # print stacktrace
        self.needsClosing=self.needsClosing-1 # unregister loop

            
    # client specific functions
    
    
    # general functions

    def send(self,packet):
        """Send a packet

        Args:
            packet (any): The packet to send. Gets encoded for you. Everything usable in an f-string (`f"{packet}"`) is valid.
        """
        self.sock.send(f"{packet}\x04".encode()) # send packet with \x04 (EOF) as end.

    def recv(self,timeout: int|None =None):
        """Receive a packet.

        Args:
            timeout (int, optional): How long to wait for a new packet. Also see Returns.

        Returns:
            any: The received packet. None if timeout is reached.
        """
        
        from time import sleep
        from select import select

        if not timeout==None: # if a timeout set
            i,tmp1,tmp2=select([self.sock], [], [], timeout) # wait for either timeout or data
            if len(i)<=0: # if timeout (select() returns the availble sockets. If len(i)<=0, no sockets are available.):
                return None # return none
        trc=bytearray()
        while True:
            rc=self.sock.recv(1) # receive 1 byte
            if not rc==b'\x04': # if byte is not \x04 (EOF)
                trc.extend(rc)  # add byte to total
            else:
                break # if byte is \x04 (EOF), stop receiving
        rc=trc.decode() # decode result
        return rc


    requestlist=[]
    
    
    
    
    # closing stuff
    
    def close(self):
        from time import sleep
        self.closing.set()
        while self.closing>0:
            sleep(.1)
            
        

