#!/usr/bin/env python3
"""
ESP32 Firmware Flash GUI Tool
Simple GUI interface for flashing firmware packages to ESP32 devices
"""

import tkinter as tk
from tkinter import ttk, filedialog, messagebox, scrolledtext
import subprocess
import sys
import threading
import zipfile
import tempfile
import os
from pathlib import Path
import serial.tools.list_ports

class ESP32FlasherGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("ESP32 Firmware Flasher")
        self.root.geometry("600x500")

        # Configure colors for dark mode compatibility
        self.setup_colors()

        # Variables
        self.port_var = tk.StringVar()
        self.firmware_path_var = tk.StringVar()
        self.temp_dir = None

        self.setup_ui()
        self.refresh_ports()

    def setup_colors(self):
        """Configure colors that work in both light and dark mode"""
        # Try to detect dark mode on macOS
        try:
            import subprocess
            result = subprocess.run(
                ['defaults', 'read', '-g', 'AppleInterfaceStyle'],
                capture_output=True,
                text=True
            )
            is_dark_mode = (result.returncode == 0 and 'Dark' in result.stdout)
        except:
            is_dark_mode = False

        if is_dark_mode:
            # Dark mode colors
            self.bg_color = '#2b2b2b'
            self.fg_color = '#ffffff'
            self.text_bg = '#1e1e1e'
            self.text_fg = '#d4d4d4'
            self.button_bg = '#3c3c3c'
        else:
            # Light mode colors
            self.bg_color = '#f0f0f0'
            self.fg_color = '#000000'
            self.text_bg = '#ffffff'
            self.text_fg = '#000000'
            self.button_bg = '#e0e0e0'

        # Configure root window
        self.root.configure(bg=self.bg_color)
        
    def setup_ui(self):
        # Main frame
        main_frame = tk.Frame(self.root, bg=self.bg_color, padx=10, pady=10)
        main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))

        # Configure grid weights
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(0, weight=1)
        main_frame.columnconfigure(1, weight=1)

        # Port selection
        tk.Label(main_frame, text="Serial Port:", bg=self.bg_color, fg=self.fg_color).grid(
            row=0, column=0, sticky=tk.W, pady=5)
        port_frame = tk.Frame(main_frame, bg=self.bg_color)
        port_frame.grid(row=0, column=1, sticky=(tk.W, tk.E), pady=5)
        port_frame.columnconfigure(0, weight=1)

        self.port_combo = ttk.Combobox(port_frame, textvariable=self.port_var, state="readonly")
        self.port_combo.grid(row=0, column=0, sticky=(tk.W, tk.E), padx=(0, 5))

        tk.Button(port_frame, text="Refresh", command=self.refresh_ports,
                  bg=self.button_bg, fg=self.fg_color, relief=tk.RAISED).grid(row=0, column=1)
        
        # Firmware package selection
        tk.Label(main_frame, text="Firmware Package:", bg=self.bg_color, fg=self.fg_color).grid(
            row=1, column=0, sticky=tk.W, pady=5)
        firmware_frame = tk.Frame(main_frame, bg=self.bg_color)
        firmware_frame.grid(row=1, column=1, sticky=(tk.W, tk.E), pady=5)
        firmware_frame.columnconfigure(0, weight=1)

        self.firmware_entry = tk.Entry(firmware_frame, textvariable=self.firmware_path_var,
                                       state="readonly", bg=self.text_bg, fg=self.text_fg,
                                       relief=tk.SUNKEN, borderwidth=2)
        self.firmware_entry.grid(row=0, column=0, sticky=(tk.W, tk.E), padx=(0, 5))

        tk.Button(firmware_frame, text="Browse", command=self.browse_firmware,
                  bg=self.button_bg, fg=self.fg_color, relief=tk.RAISED).grid(row=0, column=1)
        
        # Flash settings frame
        settings_frame = tk.LabelFrame(main_frame, text="Flash Settings",
                                       bg=self.bg_color, fg=self.fg_color,
                                       relief=tk.GROOVE, borderwidth=2, padx=5, pady=5)
        settings_frame.grid(row=2, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=10)
        settings_frame.columnconfigure(1, weight=1)

        # Flash settings
        tk.Label(settings_frame, text="Chip:", bg=self.bg_color, fg=self.fg_color).grid(
            row=0, column=0, sticky=tk.W, pady=2, padx=(0, 10))
        self.chip_var = tk.StringVar(value="esp32")
        tk.Entry(settings_frame, textvariable=self.chip_var, width=15,
                bg=self.text_bg, fg=self.text_fg, relief=tk.SUNKEN).grid(row=0, column=1, sticky=tk.W, pady=2)

        tk.Label(settings_frame, text="Baud Rate:", bg=self.bg_color, fg=self.fg_color).grid(
            row=1, column=0, sticky=tk.W, pady=2, padx=(0, 10))
        self.baud_var = tk.StringVar(value="460800")
        tk.Entry(settings_frame, textvariable=self.baud_var, width=15,
                bg=self.text_bg, fg=self.text_fg, relief=tk.SUNKEN).grid(row=1, column=1, sticky=tk.W, pady=2)

        tk.Label(settings_frame, text="Flash Mode:", bg=self.bg_color, fg=self.fg_color).grid(
            row=2, column=0, sticky=tk.W, pady=2, padx=(0, 10))
        self.flash_mode_var = tk.StringVar(value="dio")
        tk.Entry(settings_frame, textvariable=self.flash_mode_var, width=15,
                bg=self.text_bg, fg=self.text_fg, relief=tk.SUNKEN).grid(row=2, column=1, sticky=tk.W, pady=2)

        tk.Label(settings_frame, text="Flash Freq:", bg=self.bg_color, fg=self.fg_color).grid(
            row=3, column=0, sticky=tk.W, pady=2, padx=(0, 10))
        self.flash_freq_var = tk.StringVar(value="40m")
        tk.Entry(settings_frame, textvariable=self.flash_freq_var, width=15,
                bg=self.text_bg, fg=self.text_fg, relief=tk.SUNKEN).grid(row=3, column=1, sticky=tk.W, pady=2)

        tk.Label(settings_frame, text="Flash Size:", bg=self.bg_color, fg=self.fg_color).grid(
            row=4, column=0, sticky=tk.W, pady=2, padx=(0, 10))
        self.flash_size_var = tk.StringVar(value="2MB")
        tk.Entry(settings_frame, textvariable=self.flash_size_var, width=15,
                bg=self.text_bg, fg=self.text_fg, relief=tk.SUNKEN).grid(row=4, column=1, sticky=tk.W, pady=2)
        
        # Flash button
        self.flash_button = tk.Button(main_frame, text="Flash Firmware", command=self.flash_firmware,
                                      bg=self.button_bg, fg=self.fg_color,
                                      relief=tk.RAISED, padx=20, pady=5,
                                      font=('TkDefaultFont', 10, 'bold'))
        self.flash_button.grid(row=3, column=0, columnspan=2, pady=10)

        # Progress bar
        self.progress = ttk.Progressbar(main_frame, mode='indeterminate')
        self.progress.grid(row=4, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=5)

        # Log output
        tk.Label(main_frame, text="Output:", bg=self.bg_color, fg=self.fg_color).grid(
            row=5, column=0, sticky=tk.W, pady=(10, 0))
        self.log_text = scrolledtext.ScrolledText(
            main_frame,
            height=15,
            width=70,
            bg=self.text_bg,
            fg=self.text_fg,
            insertbackground=self.text_fg,  # Cursor color
            relief=tk.SUNKEN,
            borderwidth=2
        )
        self.log_text.grid(row=6, column=0, columnspan=2, sticky=(tk.W, tk.E, tk.N, tk.S), pady=5)
        main_frame.rowconfigure(6, weight=1)
        
    def refresh_ports(self):
        """Refresh the list of available serial ports"""
        ports = [port.device for port in serial.tools.list_ports.comports()]
        self.port_combo['values'] = ports
        if ports and not self.port_var.get():
            self.port_var.set(ports[0])
            
    def browse_firmware(self):
        """Browse for firmware package (.zip file)"""
        filename = filedialog.askopenfilename(
            title="Select Firmware Package",
            filetypes=[("ZIP files", "*.zip"), ("All files", "*.*")]
        )
        if filename:
            self.firmware_path_var.set(filename)
            self.validate_firmware_package(filename)
            
    def validate_firmware_package(self, zip_path):
        """Validate that the ZIP contains required firmware files"""
        required_files = [
            "bootloader.bin",
            "soundshot.bin", 
            "ota_data_initial.bin",
            "partition-table.bin"
        ]
        
        try:
            with zipfile.ZipFile(zip_path, 'r') as zip_file:
                zip_contents = zip_file.namelist()
                missing_files = []
                
                for required_file in required_files:
                    if required_file not in zip_contents:
                        missing_files.append(required_file)
                
                if missing_files:
                    messagebox.showwarning(
                        "Invalid Package", 
                        f"Missing required files:\n{', '.join(missing_files)}"
                    )
                    return False
                else:
                    self.log_message(f"✓ Valid firmware package: {Path(zip_path).name}")
                    return True
                    
        except zipfile.BadZipFile:
            messagebox.showerror("Error", "Invalid ZIP file")
            return False
            
    def log_message(self, message):
        """Add message to log output"""
        self.log_text.insert(tk.END, message + "\n")
        self.log_text.see(tk.END)
        self.root.update()
        
    def extract_firmware_package(self, zip_path):
        """Extract firmware package to temporary directory"""
        if self.temp_dir:
            # Clean up previous temp dir
            import shutil
            shutil.rmtree(self.temp_dir, ignore_errors=True)
            
        self.temp_dir = tempfile.mkdtemp(prefix="esp32_flash_")
        
        with zipfile.ZipFile(zip_path, 'r') as zip_file:
            zip_file.extractall(self.temp_dir)
            
        return self.temp_dir
        
    def flash_firmware(self):
        """Flash the firmware to ESP32"""
        # Validate inputs
        if not self.port_var.get():
            messagebox.showerror("Error", "Please select a serial port")
            return
            
        if not self.firmware_path_var.get():
            messagebox.showerror("Error", "Please select a firmware package")
            return
            
        # Disable flash button and start progress
        self.flash_button.config(state="disabled")
        self.progress.start()
        self.log_text.delete(1.0, tk.END)
        
        # Run flash in separate thread to prevent UI freezing
        thread = threading.Thread(target=self._flash_worker)
        thread.daemon = True
        thread.start()
        
    def _flash_worker(self):
        """Worker thread for flashing firmware"""
        try:
            # Extract firmware package
            self.log_message("Extracting firmware package...")
            temp_dir = self.extract_firmware_package(self.firmware_path_var.get())
            
            # Define file paths
            bootloader = os.path.join(temp_dir, "bootloader.bin")
            firmware = os.path.join(temp_dir, "soundshot.bin")
            ota_initial = os.path.join(temp_dir, "ota_data_initial.bin")
            partition = os.path.join(temp_dir, "partition-table.bin")
            
            # Build esptool command
            cmd = [
                sys.executable, "-m", "esptool",
                "--chip", self.chip_var.get(),
                "--port", self.port_var.get(),
                "--baud", self.baud_var.get(),
                "--before", "default_reset",
                "--after", "hard_reset",
                "write-flash",
                "--flash-mode", self.flash_mode_var.get(),
                "--flash-freq", self.flash_freq_var.get(),
                "--flash-size", self.flash_size_var.get(),
                "0x1000", bootloader,
                "0x20000", firmware,
                "0x11000", ota_initial,
                "0x8000", partition
            ]
            
            self.log_message(f"Running: {' '.join(cmd)}\n")
            
            # Run esptool
            process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                bufsize=1,
                universal_newlines=True
            )
            
            # Stream output
            for line in process.stdout:
                self.root.after(0, self.log_message, line.rstrip())
                
            process.wait()
            
            if process.returncode == 0:
                self.root.after(0, self.log_message, "\n✓ Flash completed successfully!")
                self.root.after(0, messagebox.showinfo, "Success", "Firmware flashed successfully!")
            else:
                self.root.after(0, self.log_message, f"\n✗ Flash failed with return code {process.returncode}")
                self.root.after(0, messagebox.showerror, "Error", "Flash operation failed!")
                
        except Exception as e:
            self.root.after(0, self.log_message, f"\n✗ Error: {str(e)}")
            self.root.after(0, messagebox.showerror, "Error", f"Flash operation failed: {str(e)}")
            
        finally:
            # Re-enable UI
            self.root.after(0, self.progress.stop)
            self.root.after(0, lambda: self.flash_button.config(state="normal"))
            
            # Clean up temp directory
            if self.temp_dir:
                import shutil
                shutil.rmtree(self.temp_dir, ignore_errors=True)
                self.temp_dir = None

def main():
    root = tk.Tk()
    app = ESP32FlasherGUI(root)
    root.mainloop()

if __name__ == "__main__":
    main()