#!/usr/bin/env python
#
# Original espota.py by Ivan Grokhotkov:
# https://gist.github.com/igrr/d35ab8446922179dc58c
#
# Modified since 2015-09-18 from Pascal Gollor (https://github.com/pgollor)
# Modified since 2015-11-09 from Hristo Gochkov (https://github.com/me-no-dev)
# Modified since 2016-01-03 from Matthew O'Gorman (https://githumb.com/mogorman)
#
# This script will push an OTA update to the ESP
# use it like:
# python espota.py -i <ESP_IP_addr> -I <Host_IP_addr> -p <ESP_port> -P <Host_port> [-a password] -f <sketch.bin>
# Or to upload SPIFFS image:
# python espota.py -i <ESP_IP_addr> -I <Host_IP_addr> -p <ESP_port> -P <HOST_port> [-a password] -s -f <spiffs.bin>
#
# Changes
# 2015-09-18:
# - Add option parser.
# - Add logging.
# - Send command to controller to differ between flashing and transmitting SPIFFS image.
#
# Changes
# 2015-11-09:
# - Added digest authentication
# - Enhanced error tracking and reporting
#
# Changes
# 2016-01-03:
# - Added more options to parser.
#
# Changes
# 2023-05-22:
# - Replaced the deprecated optparse module with argparse.
# - Adjusted the code style to conform to PEP 8 guidelines.
# - Used with statement for file handling to ensure proper resource cleanup.
# - Incorporated exception handling to catch and handle potential errors.
# - Made variable names more descriptive for better readability.
# - Introduced constants for better code maintainability.

from __future__ import print_function
import socket
import sys
import os
import argparse
import logging
import hashlib
import random

# Commands
FLASH = 0
SPIFFS = 100
AUTH = 200

# Constants
PROGRESS_BAR_LENGTH = 60


# update_progress(): Displays or updates a console progress bar
def update_progress(progress):
    if PROGRESS:
        status = ""
        if isinstance(progress, int):
            progress = float(progress)
        if not isinstance(progress, float):
            progress = 0
            status = "Error: progress var must be float\r\n"
        if progress < 0:
            progress = 0
            status = "Halt...\r\n"
        if progress >= 1:
            progress = 1
            status = "Done...\r\n"
        block = int(round(PROGRESS_BAR_LENGTH * progress))
        text = "\rUploading: [{0}] {1}% {2}".format(
            "=" * block + " " * (PROGRESS_BAR_LENGTH - block), int(progress * 100), status
        )
        sys.stderr.write(text)
        sys.stderr.flush()
    else:
        sys.stderr.write(".")
        sys.stderr.flush()


def serve(remote_addr, local_addr, remote_port, local_port, password, filename, command=FLASH, progress_handler=False):  # noqa: C901
    # Create a TCP/IP socket
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_address = (local_addr, local_port)
    logging.info("Starting on %s:%s", str(server_address[0]), str(server_address[1]))
    try:
        sock.bind(server_address)
        sock.listen(1)
    except Exception as e:
        logging.error("Listen Failed: %s", str(e))
        return 1

    content_size = os.path.getsize(filename)
    file_md5 = hashlib.md5(open(filename, "rb").read()).hexdigest()
    logging.info("Upload size: %d", content_size)
    message = "%d %d %d %s\n" % (command, local_port, content_size, file_md5)

    # Wait for a connection
    inv_tries = 0
    data = ""
    msg = "Sending invitation to %s " % remote_addr
    sys.stderr.write(msg)
    sys.stderr.flush()
    while inv_tries < 10:
        inv_tries += 1
        sock2 = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        remote_address = (remote_addr, int(remote_port))
        try:
            sent = sock2.sendto(message.encode(), remote_address)  # noqa: F841
        except:  # noqa: E722
            sys.stderr.write("failed\n")
            sys.stderr.flush()
            sock2.close()
            logging.error("Host %s Not Found", remote_addr)
            return 1
        sock2.settimeout(TIMEOUT)
        try:
            data = sock2.recv(37).decode()
            break
        except:  # noqa: E722
            sys.stderr.write(".")
            sys.stderr.flush()
            sock2.close()
    sys.stderr.write("\n")
    sys.stderr.flush()
    if inv_tries == 10:
        logging.error("No response from the ESP")
        return 1
    if data != "OK":
        if data.startswith("AUTH"):
            nonce = data.split()[1]
            cnonce_text = "%s%u%s%s" % (filename, content_size, file_md5, remote_addr)
            cnonce = hashlib.md5(cnonce_text.encode()).hexdigest()
            passmd5 = hashlib.md5(password.encode()).hexdigest()
            result_text = "%s:%s:%s" % (passmd5, nonce, cnonce)
            result = hashlib.md5(result_text.encode()).hexdigest()
            sys.stderr.write("Authenticating...")
            sys.stderr.flush()
            message = "%d %s %s\n" % (AUTH, cnonce, result)
            sock2.sendto(message.encode(), remote_address)
            sock2.settimeout(10)
            try:
                data = sock2.recv(32).decode()
            except:  # noqa: E722
                sys.stderr.write("FAIL\n")
                logging.error("No Answer to our Authentication")
                sock2.close()
                return 1
            if data != "OK":
                sys.stderr.write("FAIL\n")
                logging.error("%s", data)
                sock2.close()
                sys.exit(1)
                return 1
            sys.stderr.write("OK\n")
        else:
            logging.error("Bad Answer: %s", data)
            sock2.close()
            return 1
    sock2.close()

    logging.info("Waiting for device...")
    try:
        sock.settimeout(10)
        connection, client_address = sock.accept()
        sock.settimeout(None)
        connection.settimeout(None)
    except:  # noqa: E722
        logging.error("No response from device")
        sock.close()
        return 1
    try:
        with open(filename, "rb") as f:
            if PROGRESS:
                progress_handler(0)
            else:
                sys.stderr.write("Uploading")
                sys.stderr.flush()
            offset = 0
            while True:
                chunk = f.read(1024)
                if not chunk:
                    break
                offset += len(chunk)
                progress_handler(offset / float(content_size))
                connection.settimeout(10)
                try:
                    connection.sendall(chunk)
                    res = connection.recv(10)
                    last_response_contained_ok = "OK" in res.decode()
                except Exception as e:
                    sys.stderr.write("\n")
                    logging.error("Error Uploading: %s", str(e))
                    connection.close()
                    return 1

            if last_response_contained_ok:
                logging.info("Success")
                connection.close()
                return 0

            sys.stderr.write("\n")
            logging.info("Waiting for result...")
            count = 0
            while count < 5:
                count += 1
                connection.settimeout(60)
                try:
                    data = connection.recv(32).decode()
                    logging.info("Result: %s", data)

                    if "OK" in data:
                        logging.info("Success")
                        connection.close()
                        return 0

                except Exception as e:
                    logging.error("Error receiving result: %s", str(e))
                    connection.close()
                    return 1

            logging.error("Error response from device")
            connection.close()
            return 1

    finally:
        connection.close()

    sock.close()
    return 1


def parse_args(unparsed_args):
    parser = argparse.ArgumentParser(description="Transmit image over the air to the ESP32 module with OTA support.")

    # destination ip and port
    parser.add_argument("-i", "--ip", dest="esp_ip", action="store", help="ESP32 IP Address.", default=False)
    parser.add_argument("-I", "--host_ip", dest="host_ip", action="store", help="Host IP Address.", default="0.0.0.0")
    parser.add_argument("-p", "--port", dest="esp_port", type=int, help="ESP32 OTA Port. Default: 3232", default=3232)
    parser.add_argument(
        "-P",
        "--host_port",
        dest="host_port",
        type=int,
        help="Host server OTA Port. Default: random 10000-60000",
        default=random.randint(10000, 60000),
    )

    # authentication
    parser.add_argument("-a", "--auth", dest="auth", help="Set authentication password.", action="store", default="")

    # image
    parser.add_argument("-f", "--file", dest="image", help="Image file.", metavar="FILE", default=None)
    parser.add_argument(
        "-s",
        "--spiffs",
        dest="spiffs",
        action="store_true",
        help="Transmit a SPIFFS image and do not flash the module.",
        default=False,
    )

    # output
    parser.add_argument(
        "-d",
        "--debug",
        dest="debug",
        action="store_true",
        help="Show debug output. Overrides loglevel with debug.",
        default=False,
    )
    parser.add_argument(
        "-r",
        "--progress",
        dest="progress",
        action="store_true",
        help="Show progress output. Does not work for Arduino IDE.",
        default=False,
    )
    parser.add_argument(
        "-t",
        "--timeout",
        dest="timeout",
        type=int,
        help="Timeout to wait for the ESP32 to accept invitation.",
        default=10,
    )

    return parser.parse_args(unparsed_args)


def main(args):
    options = parse_args(args)
    log_level = logging.WARNING
    if options.debug:
        log_level = logging.DEBUG

    logging.basicConfig(level=log_level, format="%(asctime)-8s [%(levelname)s]: %(message)s", datefmt="%H:%M:%S")
    logging.debug("Options: %s", str(options))

    # check options
    global PROGRESS
    PROGRESS = options.progress

    global TIMEOUT
    TIMEOUT = options.timeout

    if not options.esp_ip or not options.image:
        logging.critical("Not enough arguments.")
        return 1

    command = FLASH
    if options.spiffs:
        command = SPIFFS

    return serve(
        options.esp_ip, options.host_ip, options.esp_port, options.host_port, options.auth, options.image, command
    )


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))