import socket
import struct
import os
import time

SERVER_IP = "192.168.1.41"
UPLOAD_LOCAL_FILE = "firmware_update.img"
UPLOAD_REMOTE_FILE = "/firmware_update.img"
DOWNLOAD_REMOTE_FILE = "/reboot.txt"
DOWNLOAD_LOCAL_FILE = "reboot.txt"

TFTP_PORT = 69
BLOCK_SIZE = 512
TIMEOUT = 3
MAX_RETRIES = 5

OP_RRQ = 1
OP_WRQ = 2
OP_DATA = 3
OP_ACK = 4
OP_ERROR = 5


def create_request(opcode, filename, mode="octet"):
    return struct.pack(f"!H{len(filename)+1}s{len(mode)+1}s",
                       opcode,
                       filename.encode() + b'\0',
                       mode.encode() + b'\0')


def create_data(block, data):
    return struct.pack("!HH", OP_DATA, block) + data


def create_ack(block):
    return struct.pack("!HH", OP_ACK, block)


def parse_packet(packet):
    opcode = struct.unpack("!H", packet[:2])[0]
    return opcode


def send_file():
    print(f"Uploading {UPLOAD_LOCAL_FILE} to {UPLOAD_REMOTE_FILE}")

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(TIMEOUT)

    try:
        # WRQ
        sock.sendto(create_request(OP_WRQ, UPLOAD_REMOTE_FILE), (SERVER_IP, TFTP_PORT))

        data, server_addr = sock.recvfrom(1024)
        if parse_packet(data) != OP_ACK:
            print("Invalid WRQ response")
            return False

        with open(UPLOAD_LOCAL_FILE, "rb") as f:
            block = 1

            while True:
                chunk = f.read(BLOCK_SIZE)
                packet = create_data(block, chunk)

                for attempt in range(MAX_RETRIES):
                    sock.sendto(packet, server_addr)
                    try:
                        ack, _ = sock.recvfrom(1024)
                        if parse_packet(ack) == OP_ACK:
                            ack_block = struct.unpack("!H", ack[2:4])[0]
                            if ack_block == block:
                                break
                    except socket.timeout:
                        if attempt == MAX_RETRIES - 1:
                            print("Upload failed (timeout)")
                            return False
                        continue

                if len(chunk) < BLOCK_SIZE:
                    print("Upload complete")
                    return True

                block = (block + 1) % 65536

    finally:
        sock.close()


def download_file():
    print(f"Triggering reset via {DOWNLOAD_REMOTE_FILE}")

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(0.5)  # short timeout

    try:
        sock.sendto(create_request(OP_RRQ, DOWNLOAD_REMOTE_FILE),
                    (SERVER_IP, TFTP_PORT))

        try:
            data, server_addr = sock.recvfrom(1024)
            opcode = parse_packet(data)

            if opcode == OP_DATA:
                block = struct.unpack("!H", data[2:4])[0]
                print(f"Received DATA block {block}, sending ACK")
                sock.sendto(create_ack(block), server_addr)

            print("Reset likely triggered (response received). Exiting.")
            return True

        except socket.timeout:
            print("No response (device likely reset). Exiting.")
            return True

    finally:
        sock.close()


if __name__ == "__main__":
    if not os.path.isfile(UPLOAD_LOCAL_FILE):
        print("Upload file not found")
    else:
        if send_file():
            time.sleep(5)  # small delay before next request
            download_file()
