"""
flash_esp8266.py

Version: 1.0.0
Author: Marcel Peterkau
Copyright: 2024
License: MIT

Beschreibung:
Dieses Skript entpackt ein ZIP-Archiv mit Firmware- und Dateisystem-Dateien und flasht diese auf einen ESP8266.
Die entpackten Dateien werden nach dem Flashen gelöscht.

Versionshistorie:
1.0.0 - Initiale Version

Benutzung:
python flash_esp8266.py <zip_file> [-v]

Optionen:
-v    Aktiviert Debug-Ausgaben
"""

import esptool
import sys
import serial.tools.list_ports
import zipfile
import os
import gzip
import shutil

BAUD_RATE = 921600  # Erhöhte Baudrate
DEBUG = '-v' in sys.argv

def debug_print(message):
    if DEBUG:
        print(message)

def get_com_ports():
    ports = [comport.device for comport in serial.tools.list_ports.comports()]
    debug_print(f"Verfügbare COM-Ports: {ports}")
    return ports

def find_new_com_port(old_ports, timeout=15):
    debug_print("Suche nach neuen COM-Ports...")
    import time
    start_time = time.time()
    while time.time() - start_time < timeout:
        new_ports = get_com_ports()
        added_ports = list(set(new_ports) - set(old_ports))
        if added_ports:
            new_port = added_ports[0]
            print(f"Neuer COM-Port gefunden: {new_port}")
            return new_port
        time.sleep(1)
    return None

def extract_files(zip_file):
    with zipfile.ZipFile(zip_file, 'r') as zip_ref:
        zip_ref.extractall()
        debug_print(f"Entpackt: {zip_ref.namelist()}")
        return zip_ref.namelist()

def decompress_gz(file_path):
    output_file = file_path.replace('.gz', '')
    with gzip.open(file_path, 'rb') as f_in:
        with open(output_file, 'wb') as f_out:
            shutil.copyfileobj(f_in, f_out)
    debug_print(f"Entpackt {file_path} zu {output_file}")
    return output_file

def clean_up(files):
    for file in files:
        if os.path.isfile(file):
            os.remove(file)
            debug_print(f"Gelöscht: {file}")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Verwendung: python flash_esp8266.py <zip_file> [-v]")
        sys.exit(1)

    ZIP_FILE = sys.argv[1]

    # Überprüfe, ob das ZIP-Archiv existiert
    if not os.path.isfile(ZIP_FILE):
        print(f"ZIP-Archiv {ZIP_FILE} nicht gefunden.")
        sys.exit(1)

    # Entpacke das ZIP-Archiv
    extracted_files = extract_files(ZIP_FILE)

    # Finde die .fw.bin- und .fs.gz-Dateien
    bin_file = None
    fs_file = None

    for file in extracted_files:
        if file.endswith('.fw.bin'):
            bin_file = file
        elif file.endswith('.fs.gz'):
            fs_file = file

    if not bin_file or not fs_file:
        print(f"Die erforderlichen Dateien (.fw.bin und .fs.gz) wurden nicht im ZIP-Archiv {ZIP_FILE} gefunden.")
        sys.exit(1)

    print(f"Gefundene Firmware-Datei: {bin_file}")
    print(f"Gefundene Dateisystem-Datei: {fs_file}")

    # Entpacke die .fs.gz-Datei
    decompressed_fs_file = decompress_gz(fs_file)

    print("Bitte stecke den ESP8266 ab, falls er angeschlossen ist, und drücke Enter.")
    input()

    print("Suche nach verfügbaren COM-Ports...")
    old_ports = get_com_ports()

    print("Bitte stecke den ESP8266 jetzt an und warte, bis der neue COM-Port erkannt wird...")

    port = find_new_com_port(old_ports, timeout=15)

    if port is None:
        print("Kein neuer COM-Port gefunden. Bitte versuche es erneut.")
        sys.exit(1)

    print(f"Neuer COM-Port gefunden: {port}")

    # Benutze esptool zum Flashen der Firmware und des Dateisystems
    try:
        # Flashen der Firmware
        esptool_args_bin = [
            '--port', port,
            '--baud', str(BAUD_RATE),
            'write_flash', '-fm', 'dout', '0x00000', bin_file
        ]
        esptool.main(esptool_args_bin)
        print("Firmware erfolgreich geflasht!")

        # Flashen des Dateisystems
        esptool_args_fs = [
            '--port', port,
            '--baud', str(BAUD_RATE),
            'write_flash', '0x300000', decompressed_fs_file
        ]
        esptool.main(esptool_args_fs)
        print("Dateisystem erfolgreich geflasht!")

        # Bereinigen der entpackten Dateien
        clean_up(extracted_files)
        clean_up([decompressed_fs_file])

    except Exception as e:
        print(f"Fehler beim Flashen: {e}")
        sys.exit(1)

    input("Drücke Enter, um das Fenster zu schließen...")  # Hält das Fenster offen