"""Stop-and-wait: a complete, working reliable protocol with a window of one. Given.

Read this before you write gbn.py. It shows every piece your protocols need:
splitting the data, the retransmission timer, what the receiver ACKs, and how
the two sides agree the transfer is over (FIN, FINACK, and the receiver's
linger). Its receiver is already a correct Go-Back-N receiver.

ACKs here are cumulative: ACK n means "I have every packet before n; send me n".
"""

import socket
import time

import packet


class Sender:
    def __init__(self, sock, peer, window=1, timeout=0.1, max_fin_tries=20):
        self.sock, self.peer = sock, peer
        self.timeout = timeout
        self.max_fin_tries = max_fin_tries
        self.stats = {"data_packets": 0, "transmissions": 0, "retransmissions": 0}

    def send(self, data):
        """Deliver `data` reliably, then close. Returns the stats dict."""
        chunks = packet.split(data)
        self.stats["data_packets"] = len(chunks)
        for seq, chunk in enumerate(chunks):
            self._until_acked(packet.encode(packet.DATA, seq, chunk), seq + 1)
        self._close(len(chunks))
        return self.stats

    def _until_acked(self, raw, want):
        first = True
        while True:
            self.sock.sendto(raw, self.peer)
            self.stats["transmissions"] += 1
            if not first:
                self.stats["retransmissions"] += 1
            first = False
            deadline = time.monotonic() + self.timeout
            while time.monotonic() < deadline:
                self.sock.settimeout(max(0.001, deadline - time.monotonic()))
                try:
                    got = packet.decode(self.sock.recvfrom(2048)[0])
                except socket.timeout:
                    break
                if got and got[0] == packet.ACK and got[1] >= want:
                    return

    def _close(self, count):
        fin = packet.encode(packet.FIN, count)
        for _ in range(self.max_fin_tries):
            self.sock.sendto(fin, self.peer)
            self.sock.settimeout(self.timeout)
            try:
                got = packet.decode(self.sock.recvfrom(2048)[0])
            except socket.timeout:
                continue
            if got and got[0] == packet.FINACK:
                return
        # every data packet was ACKed, so the data arrived; only the goodbye was lost


class Receiver:
    def __init__(self, sock, linger=0.5):
        self.sock = sock
        self.linger = linger

    def receive(self):
        """Return every byte the sender sent, in order."""
        expected, parts, peer = 0, [], None
        self.sock.settimeout(None)
        while True:
            raw, peer = self.sock.recvfrom(2048)
            got = packet.decode(raw)
            if got is None:
                continue                                   # corrupted: say nothing
            kind, seq, payload = got
            if kind == packet.DATA:
                if seq == expected:
                    parts.append(payload)
                    expected += 1
                self.sock.sendto(packet.encode(packet.ACK, expected), peer)
            elif kind == packet.FIN and seq == expected:
                self._finish(peer)
                return b"".join(parts)

    def _finish(self, peer):
        """Answer FIN, then keep answering repeats until the line goes quiet."""
        finack = packet.encode(packet.FINACK, 0)
        self.sock.sendto(finack, peer)
        self.sock.settimeout(self.linger)
        while True:
            try:
                raw, peer = self.sock.recvfrom(2048)
            except socket.timeout:
                return
            got = packet.decode(raw)
            if got and got[0] == packet.FIN:
                self.sock.sendto(finack, peer)