#!/usr/bin/env python3
"""Un TFTP mínimo (RFC 1350), solo lectura: un protocolo real de stop-and-wait sobre UDP.

    python3 tftp_mini.py serve <carpeta> [--timeout-ms 200] [--port 69]
    python3 tftp_mini.py get <ip> <archivo> <destino> [--port 69] [--ignore-block-number]

Por cada bloque de 512 bytes el servidor espera su ACK. Si el temporizador vence, repite el
bloque. Un bloque de menos de 512 bytes indica el final. --ignore-block-number hace que el
cliente escriba todo lo que llega, sin mirar el número de bloque: sirve para ver qué pasa
sin números de secuencia.
"""
import argparse
import hashlib
import os
import socket
import struct
import threading
import time

OP_RRQ, OP_DATA, OP_ACK, OP_ERROR = 1, 3, 4, 5
BLOCK_SIZE = 512
MAX_TRIES = 40
GET_IDLE_S = 10


def send_file(path, client, timeout_s):
    """Atiende una lectura, desde un socket y un puerto nuevos, como pide el RFC 1350."""
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.connect(client)
    block_number = 0
    retransmissions = 0
    try:
        with open(path, "rb") as file:
            while True:
                data = file.read(BLOCK_SIZE)
                block_number += 1
                packet = struct.pack("!HH", OP_DATA, block_number & 0xffff) + data
                retransmissions += stop_and_wait(sock, packet, block_number & 0xffff, timeout_s)
                if len(data) < BLOCK_SIZE:
                    break
        print(f"enviado {os.path.basename(path)}: {block_number} bloques, "
              f"{retransmissions} retransmisiones", flush=True)
    except FileNotFoundError:
        sock.send(struct.pack("!HH", OP_ERROR, 1) + b"no existe\0")
    except (TimeoutError, ConnectionError) as error:
        print(f"transferencia abandonada: {error}", flush=True)
    finally:
        sock.close()


def stop_and_wait(sock, packet, block_number, timeout_s):
    """Manda un bloque y espera su ACK. Devuelve cuántas veces tuvo que repetirlo."""
    for tries in range(MAX_TRIES):
        sock.send(packet)
        deadline = time.monotonic() + timeout_s
        while (left := deadline - time.monotonic()) > 0:
            sock.settimeout(left)
            try:
                reply = sock.recv(64)
            except TimeoutError:
                break
            # Un ACK viejo se ignora: repetir por eso multiplicaría los bloques en la red.
            if len(reply) == 4 and struct.unpack("!HH", reply) == (OP_ACK, block_number):
                return tries
    raise TimeoutError(f"sin ACK del bloque {block_number}")


def serve(folder, port, timeout_ms):
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(("0.0.0.0", port))
    print(f"TFTP en el puerto {port}, carpeta {folder}", flush=True)
    while True:
        request, client = sock.recvfrom(1024)
        if len(request) < 4 or struct.unpack("!H", request[:2])[0] != OP_RRQ:
            continue
        name = request[2:].split(b"\0")[0].decode(errors="replace")
        path = os.path.join(folder, os.path.basename(name))
        threading.Thread(target=send_file, args=(path, client, timeout_ms / 1000),
                         daemon=True).start()


def get(server_ip, name, target, port, ignore_block_number):
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(1.0)
    request = struct.pack("!H", OP_RRQ) + name.encode() + b"\0octet\0"
    digest = hashlib.sha256()
    expected = 1
    duplicates = 0
    total = 0
    started = time.monotonic()
    last_heard = started
    sock.sendto(request, (server_ip, port))
    with open(target, "wb") as file:
        while True:
            try:
                packet, sender = sock.recvfrom(BLOCK_SIZE + 4)
            except TimeoutError:
                if time.monotonic() - last_heard > GET_IDLE_S:
                    raise SystemExit("el servidor dejó de responder") from None
                if expected == 1:
                    sock.sendto(request, (server_ip, port))
                continue
            last_heard = time.monotonic()
            opcode, block = struct.unpack("!HH", packet[:4])
            if opcode == OP_ERROR:
                raise SystemExit(f"error del servidor: {packet[4:-1].decode(errors='replace')}")
            if opcode != OP_DATA:
                continue
            is_new = block == (expected & 0xffff)
            sock.sendto(struct.pack("!HH", OP_ACK, block), sender)
            if not is_new:
                duplicates += 1
            if is_new or ignore_block_number:
                data = packet[4:]
                file.write(data)
                digest.update(data)
                total += len(data)
            if is_new:
                expected += 1
                if len(packet) - 4 < BLOCK_SIZE:
                    break
    elapsed = time.monotonic() - started
    rate = total * 8 / elapsed / 1000
    print(f"bytes={total} tiempo={elapsed:.2f}s tasa={rate:.1f} kbit/s "
          f"duplicados={duplicates} sha256={digest.hexdigest()[:16]}")


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    commands = parser.add_subparsers(dest="command", required=True)
    server = commands.add_parser("serve")
    server.add_argument("folder")
    server.add_argument("--port", type=int, default=69)
    server.add_argument("--timeout-ms", type=int, default=200)
    client = commands.add_parser("get")
    client.add_argument("server_ip")
    client.add_argument("name")
    client.add_argument("target")
    client.add_argument("--port", type=int, default=69)
    client.add_argument("--ignore-block-number", action="store_true")
    args = parser.parse_args()
    if args.command == "serve":
        serve(args.folder, args.port, args.timeout_ms)
    else:
        get(args.server_ip, args.name, args.target, args.port, args.ignore_block_number)


if __name__ == "__main__":
    main()