# 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 re
import subprocess
import sys
from typing import Any, Dict, Optional

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


class ChipInfoManager:
    """
    Manages retrieval of chip information using esptool and espefuse.
    """

    def __init__(self, port: str, baudrate: int = 115200):
        self.port = port
        self.baudrate = baudrate

    def _run_tool(self, tool_name: str, args: list, progress_callback=None) -> str:
        # Try various methods to run the tool
        methods = [
            [sys.executable, "-m", tool_name],
            [tool_name],
            [f"{tool_name}.py"],
        ]

        last_error = ""
        for cmd_base in methods:
            cmd = cmd_base + ["--port", self.port, "--baud", str(self.baudrate)] + args
            try:
                if progress_callback:
                    return self._run_tool_interactive(cmd, progress_callback)
                
                result = subprocess.run(
                    cmd, capture_output=True, text=True, check=False
                )
                if result.returncode == 0:
                    return result.stdout
                else:
                    last_error = result.stderr
            except Exception as e:
                last_error = str(e)

        logger.error(f"All methods failed to run {tool_name}. Last error: {last_error}")
        return ""

    def _run_tool_interactive(self, cmd: list, progress_callback) -> str:
        """Runs tool and parses progress from stdout/stderr."""
        full_output = []
        try:
            process = subprocess.Popen(
                cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1
            )
            
            # Progress regex patterns for esptool/espefuse
            # e.g. "40960/1048576 bytes... (3%)" or "100.0%"
            progress_re = re.compile(r"(\d+\.?\d*)\s*%")
            
            for line in process.stdout:
                full_output.append(line)
                
                # Try to find percentage
                match = progress_re.search(line)
                if match:
                    try:
                        percent = float(match.group(1))
                        progress_callback(percent)
                    except:
                        pass
            
            process.wait()
            return "".join(full_output)
        except Exception as e:
            logger.error(f"Failed to run tool interactive: {e}")
            return ""

    def get_all_info(self) -> Dict[str, Any]:
        """Gathers all chip, flash, adc, efuse and security info."""
        info = {
            "chip": self._get_chip_summary(),
            "flash": self._get_flash_summary(),
            "adc": self._get_adc_summary(),
            "efuse": self._get_efuse_summary(),
            "security": self._get_security_summary(),
        }
        return info

    def _clean_ansi(self, text: str) -> str:
        ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
        return ansi_escape.sub("", text)

    def _get_chip_summary(self) -> Dict[str, Any]:
        stdout = self._run_tool("esptool", ["chip_id"])
        stdout_clean = self._clean_ansi(stdout)
        summary = {"raw": stdout_clean}

        # Try multiple patterns for chip model
        model_patterns = [
            r"Chip is (.*?)(?: \(revision (.*?)\))?[\s\r\n]",
            r"Chip type:\s+(.*)",
            r"Detecting chip type\.\.\.\s+(.*)",
        ]

        for pattern in model_patterns:
            match = re.search(pattern, stdout_clean)
            if match:
                summary["model"] = match.group(1).strip()
                if len(match.groups()) > 1 and match.group(2):
                    summary["revision"] = match.group(2).strip()
                break

        if "model" not in summary:
            # Fallback for P4 or newer tools
            match_generic = re.search(r"ESP32-[PQ]\d+", stdout_clean)
            if match_generic:
                summary["model"] = match_generic.group(0)

        match_mac = re.search(r"MAC: ([\da-fA-F:]+)", stdout_clean)
        if match_mac:
            summary["mac"] = match_mac.group(1).strip()

        match_features = re.search(r"Features: (.*)", stdout_clean)
        if match_features:
            summary["features"] = match_features.group(1).strip()

        return summary

    def _get_flash_summary(self) -> Dict[str, Any]:
        stdout = self._run_tool("esptool", ["flash_id"])
        stdout_clean = self._clean_ansi(stdout)
        summary = {"raw": stdout_clean}

        match_man = re.search(r"Manufacturer: (\w+)", stdout_clean)
        if match_man:
            summary["manufacturer_id"] = match_man.group(1)

        match_dev = re.search(r"Device: (\w+)", stdout_clean)
        if match_dev:
            summary["device_id"] = match_dev.group(1)

        match_size = re.search(r"Detected flash size: (.*)", stdout_clean)
        if match_size:
            summary["size"] = match_size.group(1).strip()

        # Optional: status register
        stdout_status = self._run_tool("esptool", ["read_flash_status"])
        match_status = re.search(r"Status value: (0x[\da-fA-F]+)", stdout_status)
        if match_status:
            summary["status"] = match_status.group(1)

        return summary

    def _get_adc_summary(self) -> str:
        stdout = self._run_tool("espefuse", ["adc_info"])
        return self._clean_ansi(stdout).strip()

    def _get_efuse_summary(self) -> str:
        stdout = self._run_tool("espefuse", ["summary"])
        return self._clean_ansi(stdout).strip()

    def _get_security_summary(self) -> Dict[str, Any]:
        stdout = self._run_tool("esptool", ["get_security_info"])
        stdout_clean = self._clean_ansi(stdout)
        summary = {"raw": stdout_clean}

        # Simple parsing for quick view
        if "Flash Encryption: Enabled" in stdout_clean:
            summary["flash_encryption"] = "Enabled"
        else:
            summary["flash_encryption"] = "Disabled"

        if "Secure Boot: Enabled" in stdout_clean:
            summary["secure_boot"] = "Enabled"
        else:
            summary["secure_boot"] = "Disabled"

        return summary

    def get_partition_table(self) -> Dict[str, Any]:
        """Reads and parses the partition table from flash."""
        try:
            import os
            import struct
            import tempfile

            with tempfile.NamedTemporaryFile(delete=False) as tmp:
                tmp_name = tmp.name

            # typical pt offset is 0x8000, read 4KB
            self._run_tool("esptool", ["read_flash", "0x8000", "0x1000", tmp_name])

            partitions = []
            with open(tmp_name, "rb") as f:
                data = f.read()

            os.remove(tmp_name)

            # Parse partition table binary format
            # Each entry is 32 bytes: magic(2), type(1), subtype(1), offset(4), size(4), label(16), flags(4)
            for i in range(0, len(data), 32):
                entry = data[i : i + 32]
                if len(entry) < 32:
                    break
                magic = struct.unpack("<H", entry[0:2])[0]
                if magic != 0x50AA:
                    if magic == 0xFFFF:  # End of table
                        break
                    continue  # Skip invalid

                ptype, subtype, offset, size = struct.unpack("<BBII", entry[2:12])
                label = entry[12:28].split(b"\x00")[0].decode("utf-8", "ignore")
                flags = struct.unpack("<I", entry[28:32])[0]

                type_str = (
                    "app" if ptype == 0x00 else "data" if ptype == 0x01 else hex(ptype)
                )

                partitions.append(
                    {
                        "label": label,
                        "type": type_str,
                        "subtype": hex(subtype),
                        "offset": hex(offset),
                        "size": hex(size),
                        "size_bytes": size,
                        "flags": hex(flags),
                    }
                )

            return {"status": "success", "partitions": partitions}
        except Exception as e:
            return {"status": "error", "error": str(e)}

    def get_app_info(self, offset: str, size: str = "0x100000", progress_callback=None) -> Dict[str, Any]:
        """Reads app partition and uses image_info to parse it."""
        try:
            import os
            import tempfile

            with tempfile.NamedTemporaryFile(delete=False, suffix=".bin") as tmp:
                tmp_name = tmp.name

            # Read full app partition for absolute accuracy
            self._run_tool("esptool", ["read_flash", offset, size, tmp_name], progress_callback=progress_callback)

            # Try image-info first, then fallback to image_info
            stdout = self._run_tool("esptool", ["image-info", tmp_name])
            if not stdout or "A fatal error occurred" in stdout:
                stdout = self._run_tool("esptool", ["image_info", tmp_name])
            
            # We DON'T remove it immediately if we want standard download? 
            # Actually backend will handle moving it to Data/dumps if needed.
            # But for App Info, we just want the metadata.
            # However, the user said "保存成一个 .bin 一个 .json，默认就放到前端"
            # So maybe we should save it to a persistent place.
            
            # For now, extract metadata
            stdout_clean = self._clean_ansi(stdout)
            info = {"raw": stdout_clean, "status": "success", "tmp_path": tmp_name}

            match_name = re.search(r"Project name:\s+(.*)", stdout_clean)
            if match_name:
                info["project_name"] = match_name.group(1).strip()

            match_ver = re.search(r"App version:\s+(.*)", stdout_clean)
            if match_ver:
                info["app_version"] = match_ver.group(1).strip()

            match_time = re.search(r"Compile time:\s+(.*)", stdout_clean)
            if match_time:
                info["compile_time"] = match_time.group(1).strip()

            return info
        except Exception as e:
            return {"status": "error", "error": str(e)}

    def read_flash(self, offset: str, size: str, output_path: str, progress_callback=None) -> Dict[str, Any]:
        """Reads a range of flash and saves it to a file."""
        try:
            stdout = self._run_tool("esptool", ["read_flash", offset, size, output_path], progress_callback=progress_callback)
            if os.path.exists(output_path):
                return {"status": "success", "path": output_path, "raw": stdout}
            return {"status": "error", "error": "Output file not created", "raw": stdout}
        except Exception as e:
            return {"status": "error", "error": str(e)}
