# Copyright 2026 Espressif Systems (Shanghai) Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import os
import subprocess
import threading
import time
from typing import Optional

from esp_viewer.core.protocol import EventType, LogLevel, Packet

logger = logging.getLogger("esp_viewer.firmware")


class FirmwareManager:
    """
    Manages firmware downloading (from GitHub Releases/etc) and
    flashing to the ESP device using esptool.py in subprocesses.
    """

    def __init__(self, engine):
        # We need a reference to the Engine to emit progress updates and log messages
        # back to the generic UI via event queues.
        self.engine = engine
        self.download_thread: Optional[threading.Thread] = None
        self.flash_thread: Optional[threading.Thread] = None
        self._cache_dir = "Data/firmware"
        self.latest_firmware_path: Optional[str] = None
        os.makedirs(self._cache_dir, exist_ok=True)

    def _log(self, level: LogLevel, msg: str):
        """Helper to pipe firmware logs into the engine's main system log."""
        self.engine.event_queue.put(
            Packet(
                EventType.LOG_MSG, {"level": level.value, "msg": f"[Firmware] {msg}"}
            )
        )
        if level == LogLevel.DEBUG:
            logger.debug(msg)
        elif level == LogLevel.INFO:
            logger.info(msg)
        elif level == LogLevel.WARNING:
            logger.warning(msg)
        elif level == LogLevel.ERROR:
            logger.error(msg)

    def _emit_progress(self, phase: str, percent: float, msg: str = ""):
        self.engine.event_queue.put(
            Packet(
                EventType.FIRMWARE_PROGRESS,
                {"phase": phase, "percent": percent, "message": msg},
            )
        )

    def start_download(self, url: str):
        """
        Background thread trigger. Downloads firmware from `url`.
        If `url == 'test'`, generates an all-zero 1MB mock binary.
        """
        if self.download_thread and self.download_thread.is_alive():
            self._log(LogLevel.WARNING, "Download already in progress.")
            return

        def _download_task():
            if url == "test":
                self._log(LogLevel.INFO, "Starting Mock Firmware Download...")
                self._emit_progress("download", 0, "Initializing mock download")
                mock_path = os.path.join(self._cache_dir, "mock_firmware_1MB.bin")
                try:
                    with open(mock_path, "wb") as f:
                        total_chunks = 100
                        chunk_size = 1024 * 10
                        for i in range(total_chunks):
                            time.sleep(0.02)
                            f.write(b"\x00" * chunk_size)
                            self._emit_progress(
                                "download",
                                ((i + 1) / total_chunks) * 100,
                                f"Downloading chunk {i+1}/{total_chunks}",
                            )
                    self._emit_progress("download", 100, "Download Complete")
                    self._log(LogLevel.INFO, f"Mock firmware saved to {mock_path}")
                    self.latest_firmware_path = mock_path
                except Exception as e:
                    self._emit_progress("download", 0, "Download Failed")
                    self._log(LogLevel.ERROR, f"Download Exception: {e}", exc_info=True)
            else:
                self._log(LogLevel.INFO, f"Starting firmware download from: {url}")
                self._emit_progress("download", 0, "Connecting to server...")
                import ssl
                import urllib.error
                import urllib.request

                try:
                    # Extract filename from URL or default
                    filename = (
                        url.split("/")[-1]
                        if "/" in url and not url.endswith("/")
                        else "firmware.bin"
                    )
                    if "?" in filename:
                        filename = filename.split("?")[0]
                    if not filename.endswith((".bin", ".zip")):
                        filename += ".bin"
                    download_path = os.path.join(self._cache_dir, filename)

                    # Unverified context for development convenience vs HTTPS issues
                    ctx = ssl.create_default_context()
                    ctx.check_hostname = False
                    ctx.verify_mode = ssl.CERT_NONE

                    def report_hook(count, block_size, total_size):
                        if total_size > 0:
                            percent = min(
                                100.0, (count * block_size * 100.0) / total_size
                            )
                            self._emit_progress(
                                "download", percent, f"Downloading: {percent:.1f}%"
                            )

                    with urllib.request.urlopen(url, context=ctx) as response, open(
                        download_path, "wb"
                    ) as out_file:
                        total_size = int(response.info().get("Content-Length", -1))
                        count = 0
                        block_size = 1024 * 8
                        while True:
                            buffer = response.read(block_size)
                            if not buffer:
                                break
                            count += 1
                            out_file.write(buffer)
                            report_hook(count, block_size, total_size)

                    self._emit_progress("download", 100, "Download Complete")
                    self._log(LogLevel.INFO, f"Firmware saved to {download_path}")
                    self.latest_firmware_path = download_path
                except urllib.error.URLError as e:
                    self._emit_progress("download", 0, f"Network Error: {e.reason}")
                    self._log(LogLevel.ERROR, f"Download URLError: {e.reason}")
                except Exception as e:
                    self._emit_progress("download", 0, "Download Failed")
                    self._log(LogLevel.ERROR, f"Download Exception: {e}", exc_info=True)

        self.download_thread = threading.Thread(target=_download_task, daemon=True)
        self.download_thread.start()

    def flash_firmware(self, port: str, firmware_path: Optional[str] = None):
        """
        Invokes esptool.py inside a subprocess to flash the
        previously downloaded (or mock) binary.
        """
        if self.flash_thread and self.flash_thread.is_alive():
            self._log(LogLevel.WARNING, "Flashing already in progress.")
            return

        # 1. Inform Engine to release the serial port so esptool can claim it.
        # This prevents "Access is denied" or "Device Busy" errors.
        self.engine.device_service._disconnect()

        def _flash_task():
            self._log(LogLevel.INFO, f"Preparing to flash device on port: {port}")
            self._emit_progress("flash", 0, "Initializing esptool")

            path_to_flash = firmware_path or self.latest_firmware_path

            if not path_to_flash or not os.path.exists(path_to_flash):
                self._log(
                    LogLevel.ERROR,
                    f"Firmware binary not found at '{path_to_flash}'. Please download first.",
                )
                self._emit_progress("flash", 0, "Missing firmware file")
                return

            # Note: For real flashing we'd specify specific layout addresses:
            # esptool.py --port /dev/ttyUSB0 --baud 460800 write_flash 0x0 bootloader.bin 0x8000 partition.bin 0x10000 app.bin
            # Here we just flash the 1MB mock to 0x10000 as a proof of concept.

            import sys

            cmd = [
                sys.executable,
                "-m",
                "esptool",
                "--port",
                port,
                "--baud",
                "460800",
                "write_flash",
                "-z",
                "0x10000",
                path_to_flash,
            ]

            self._log(LogLevel.DEBUG, f"Executing: {' '.join(cmd)}")

            try:
                # Use sub-process to run esptool asynchronously
                process = subprocess.Popen(
                    cmd,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    text=True,
                    bufsize=1,  # Line buffered
                    universal_newlines=True,
                )

                # Parse stdout line-by-line in real-time
                for line in iter(process.stdout.readline, ""):
                    line = line.strip()
                    if not line:
                        continue

                    # Pipe esptool output directly to our System Log window
                    self._log(LogLevel.INFO, f"[esptool] {line}")

                    # Optional: Parse percentage from lines like "Writing at 0x00018000... (16 %)"
                    if "Writing at " in line and "%" in line:
                        try:
                            # Crude extraction of percentage
                            pct_str = line.split("(")[-1].split("%")[0].strip()
                            pct = float(pct_str)
                            self._emit_progress("flash", pct, line)
                        except:
                            pass

                process.stdout.close()
                return_code = process.wait()

                if return_code == 0:
                    self._emit_progress(
                        "flash", 100, "Flashing Complete. Please reset device."
                    )
                    self._log(LogLevel.INFO, "Flashing successfully completed.")
                else:
                    self._emit_progress(
                        "flash", 0, f"Flashing failed with code {return_code}"
                    )
                    self._log(
                        LogLevel.ERROR, f"Flashing failed with code {return_code}"
                    )

            except Exception as e:
                self._emit_progress("flash", 0, "Flashing exception")
                self._log(LogLevel.ERROR, f"Flashing Exception: {e}", exc_info=True)
            finally:
                # Re-engage the serial port for telemetry
                self.engine.device_service.connected = False
                # Optionally trigger auto-reconnect here if the user's config requires it

        self.flash_thread = threading.Thread(target=_flash_task, daemon=True)
        self.flash_thread.start()
