import socket
import threading
import protocol

LOCAL_PORT = 9000
TIMEOUT = 2

class StopWaitChat:
    def __init__(self, local_port, peer_addr):
        self.sock = socket.socket(
            socket.AF_INET6,
            socket.SOCK_DGRAM
        )
        self.sock.bind(("::", local_port))
        self.sock.settimeout(TIMEOUT)

        self.peer_addr = peer_addr
        self.send_seq = 0
        self.recv_seq = 0

    def sender(self):
        while True:
            msg = input(">> ").encode()
            pkt = protocol.make_packet(
                self.send_seq,
                protocol.TYPE_DATA,
                msg
            )

            while True:
                self.sock.sendto(pkt, self.peer_addr)
                try:
                    data, _ = self.sock.recvfrom(1024)
                    seq, ptype, _ = protocol.parse_packet(data)
                    if ptype == protocol.TYPE_ACK and seq == self.send_seq:
                        self.send_seq ^= 1
                        break
                except socket.timeout:
                    print("[timeout] retransmit")

    def receiver(self):
        while True:
            data, addr = self.sock.recvfrom(1024)
            seq, ptype, payload = protocol.parse_packet(data)

            if ptype == protocol.TYPE_DATA:
                if seq == self.recv_seq:
                    print("\npeer:", payload.decode())
                    self.recv_seq ^= 1
                ack = protocol.make_packet(seq, protocol.TYPE_ACK)
                self.sock.sendto(ack, addr)

if __name__ == "__main__":
    import sys
    if len(sys.argv) != 3:
        print("用法: python3 chat.py <对方IPv6地址> <对方端口>")
        exit(1)

    peer_ip = sys.argv[1]
    peer_port = int(sys.argv[2])

    chat = StopWaitChat(
        LOCAL_PORT,
        (peer_ip, peer_port)
    )

    threading.Thread(target=chat.receiver, daemon=True).start()
    chat.sender()