"""
Flash Manager Component

Provides ESP flash memory operations: write, read, erase, and backup.
All operations shell out to esptool as an async subprocess, matching
the pattern established in chip_control.py.
"""

import asyncio
import logging
import re
import time
from pathlib import Path
from typing import Any

from fastmcp import Context, FastMCP

from ..config import ESPToolServerConfig

logger = logging.getLogger(__name__)


class FlashManager:
    """ESP flash memory management and operations"""

    def __init__(self, app: FastMCP, config: ESPToolServerConfig) -> None:
        self.app = app
        self.config = config
        self._register_tools()

    async def _run_esptool(
        self,
        port: str,
        args: list[str],
        timeout: float = 120.0,
    ) -> dict[str, Any]:
        """Run esptool with arbitrary args as an async subprocess.

        Args:
            port: Serial port or socket:// URI
            args: esptool arguments after --port (e.g. ["write-flash", "0x0", "fw.bin"])
            timeout: Timeout in seconds (flash operations can be slow)

        Returns:
            dict with "success", "output", and optionally "error"
        """
        cmd = [
            self.config.esptool_path,
            "--port", port,
            *args,
        ]
        proc = None
        try:
            proc = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
            stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
            output = (stdout or b"").decode() + (stderr or b"").decode()

            if proc.returncode != 0:
                return {"success": False, "error": output.strip()}

            return {"success": True, "output": output}

        except asyncio.TimeoutError:
            if proc and proc.returncode is None:
                proc.kill()
                await proc.wait()
            return {"success": False, "error": f"Timeout after {timeout}s"}
        except FileNotFoundError:
            return {
                "success": False,
                "error": f"esptool not found at {self.config.esptool_path}",
            }
        except Exception as e:
            if proc and proc.returncode is None:
                proc.kill()
                await proc.wait()
            return {"success": False, "error": str(e)}

    def _register_tools(self) -> None:
        """Register flash management tools"""

        @self.app.tool("esp_flash_firmware")
        async def flash_firmware(
            context: Context,
            firmware_path: str,
            port: str | None = None,
            address: str = "0x0",
            verify: bool = True,
        ) -> dict[str, Any]:
            """Flash firmware to ESP device.

            Writes a binary firmware file to the device's flash memory using esptool.
            Supports any port including socket:// URIs for QEMU virtual devices.

            Args:
                firmware_path: Path to the firmware binary (.bin) to flash
                port: Serial port or socket:// URI (auto-detect if not specified)
                address: Flash address to write to (hex string, default: "0x0").
                    Use partition offsets for non-firmware images (e.g. "0x290000" for LittleFS).
                verify: Verify flash contents after writing (default: true)
            """
            return await self._flash_firmware_impl(context, firmware_path, port, address, verify)

        @self.app.tool("esp_flash_read")
        async def flash_read(
            context: Context,
            output_path: str,
            port: str | None = None,
            start_address: str = "0x0",
            size: str | None = None,
        ) -> dict[str, Any]:
            """Read flash memory contents to a file.

            Reads raw bytes from flash and saves to the specified output path.
            If size is not specified, reads the entire flash.

            Args:
                output_path: File path to save the flash contents
                port: Serial port or socket:// URI (auto-detect if not specified)
                start_address: Flash offset to start reading from (hex string, default: "0x0")
                size: Number of bytes to read (hex or decimal string, reads all if not specified)
            """
            return await self._flash_read_impl(context, output_path, port, start_address, size)

        @self.app.tool("esp_flash_erase")
        async def flash_erase(
            context: Context,
            port: str | None = None,
            start_address: str = "0x0",
            size: str | None = None,
        ) -> dict[str, Any]:
            """Erase flash memory regions.

            Erases the entire flash if no start_address and size are given.
            Otherwise erases the specified region. Erased bytes become 0xFF.

            Args:
                port: Serial port or socket:// URI (auto-detect if not specified)
                start_address: Flash offset to start erasing (hex string, default: "0x0")
                size: Number of bytes to erase (hex or decimal string, erases all if not specified)
            """
            return await self._flash_erase_impl(context, port, start_address, size)

        @self.app.tool("esp_flash_backup")
        async def flash_backup(
            context: Context,
            backup_path: str,
            port: str | None = None,
            include_bootloader: bool = True,
        ) -> dict[str, Any]:
            """Create complete flash backup to a file.

            Reads the entire flash contents and saves to the specified path.
            The resulting file can be restored with esp_flash_firmware.

            Args:
                backup_path: File path to save the flash backup
                port: Serial port or socket:// URI (auto-detect if not specified)
                include_bootloader: Start from address 0x0 to include bootloader (default: true)
            """
            return await self._flash_backup_impl(context, backup_path, port, include_bootloader)

        @self.app.tool("esp_flash_multi")
        async def flash_multi(
            context: Context,
            files: list[dict],
            port: str | None = None,
            verify: bool = True,
            compress: bool = True,
        ) -> dict[str, Any]:
            """Flash multiple binaries at different addresses in one operation.

            Writes multiple binary files to flash in a single esptool invocation.
            Faster than multiple separate flash operations since it connects once.

            Common use cases:
            - Flash bootloader + partition table + app in one shot
            - Deploy complete firmware stack with filesystem image

            Args:
                files: List of dicts with "address" (hex string) and "path" (file path).
                    Example: [{"address": "0x0", "path": "bootloader.bin"},
                              {"address": "0x8000", "path": "partitions.bin"},
                              {"address": "0x10000", "path": "app.bin"}]
                port: Serial port or socket:// URI (required)
                verify: Verify flash contents after writing (default: true)
                compress: Use compression for faster transfer (default: true)
            """
            return await self._flash_multi_impl(context, files, port, verify, compress)

        @self.app.tool("esp_verify_flash")
        async def verify_flash(
            context: Context,
            firmware_path: str,
            port: str | None = None,
            address: str = "0x0",
        ) -> dict[str, Any]:
            """Verify flash contents match a file without re-flashing.

            Reads flash memory at the specified address and compares against
            the provided file. Useful for confirming successful flash operations
            or checking if an update is needed.

            Args:
                firmware_path: Path to the binary file to compare against
                port: Serial port or socket:// URI (required)
                address: Flash address to verify from (hex string, default: "0x0")
            """
            return await self._verify_flash_impl(context, firmware_path, port, address)

    async def _flash_firmware_impl(
        self,
        context: Context,
        firmware_path: str,
        port: str | None,
        address: str,
        verify: bool,
    ) -> dict[str, Any]:
        """Write firmware to flash via esptool write-flash."""

        fw_path = Path(firmware_path)
        if not fw_path.exists():
            return {"success": False, "error": f"Firmware file not found: {firmware_path}"}

        if not port:
            return {"success": False, "error": "Port is required (no auto-detect for flash operations)"}

        start_time = time.time()

        args = ["write-flash", address, str(fw_path)]
        if not verify:
            args.insert(0, "--no-verify")

        result = await self._run_esptool(port, args, timeout=180.0)

        if not result["success"]:
            return {
                "success": False,
                "error": result["error"],
                "port": port,
                "firmware_path": firmware_path,
                "address": address,
            }

        output = result["output"]
        elapsed = round(time.time() - start_time, 1)

        # Parse bytes written from output
        bytes_written = 0
        write_matches = re.findall(r"Wrote (\d+) bytes", output)
        for match in write_matches:
            bytes_written += int(match)

        verified = "Hash of data verified" in output or "Verified" in output

        return {
            "success": True,
            "port": port,
            "firmware_path": firmware_path,
            "address": address,
            "firmware_size": fw_path.stat().st_size,
            "bytes_written": bytes_written,
            "verified": verified if verify else None,
            "elapsed_seconds": elapsed,
        }

    async def _flash_read_impl(
        self,
        context: Context,
        output_path: str,
        port: str | None,
        start_address: str,
        size: str | None,
    ) -> dict[str, Any]:
        """Read flash contents via esptool read-flash."""

        if not port:
            return {"success": False, "error": "Port is required (no auto-detect for flash operations)"}

        # Determine read size — if not specified, read entire flash (detect first)
        if not size:
            detect = await self._run_esptool(port, ["flash-id"], timeout=15.0)
            if not detect["success"]:
                return {"success": False, "error": f"Could not detect flash size: {detect['error']}"}

            # Parse flash size from output
            flash_size_match = re.search(r"Detected flash size:\s*(\d+)([KMG]B)", detect["output"])
            if flash_size_match:
                num = int(flash_size_match.group(1))
                unit = flash_size_match.group(2)
                multiplier = {"KB": 1024, "MB": 1024 * 1024, "GB": 1024 * 1024 * 1024}
                size = str(num * multiplier.get(unit, 1))
            else:
                return {"success": False, "error": "Could not determine flash size. Specify size manually."}

        # Ensure output directory exists
        out = Path(output_path)
        out.parent.mkdir(parents=True, exist_ok=True)

        start_time = time.time()
        result = await self._run_esptool(
            port,
            ["read-flash", start_address, size, str(out)],
            timeout=300.0,
        )

        if not result["success"]:
            return {"success": False, "error": result["error"], "port": port}

        elapsed = round(time.time() - start_time, 1)

        return {
            "success": True,
            "port": port,
            "output_path": str(out),
            "start_address": start_address,
            "bytes_read": out.stat().st_size if out.exists() else 0,
            "elapsed_seconds": elapsed,
        }

    async def _flash_erase_impl(
        self,
        context: Context,
        port: str | None,
        start_address: str,
        size: str | None,
    ) -> dict[str, Any]:
        """Erase flash via esptool erase-flash or erase-region."""

        if not port:
            return {"success": False, "error": "Port is required (no auto-detect for flash operations)"}

        start_time = time.time()

        if size:
            # Erase specific region
            result = await self._run_esptool(
                port,
                ["erase-region", start_address, size],
                timeout=60.0,
            )
        else:
            # Erase entire flash
            result = await self._run_esptool(
                port,
                ["erase-flash"],
                timeout=60.0,
            )

        if not result["success"]:
            return {"success": False, "error": result["error"], "port": port}

        elapsed = round(time.time() - start_time, 1)

        return {
            "success": True,
            "port": port,
            "erase_type": "region" if size else "full",
            "start_address": start_address if size else "0x0",
            "size": size,
            "elapsed_seconds": elapsed,
        }

    async def _flash_backup_impl(
        self,
        context: Context,
        backup_path: str,
        port: str | None,
        include_bootloader: bool,
    ) -> dict[str, Any]:
        """Read entire flash to create a backup file."""

        start_address = "0x0" if include_bootloader else "0x1000"

        return await self._flash_read_impl(
            context,
            output_path=backup_path,
            port=port,
            start_address=start_address,
            size=None,  # auto-detect full flash
        )

    async def _flash_multi_impl(
        self,
        context: Context,
        files: list[dict],
        port: str | None,
        verify: bool,
        compress: bool,
    ) -> dict[str, Any]:
        """Flash multiple binaries at different addresses via esptool write-flash."""

        if not port:
            return {"success": False, "error": "Port is required (no auto-detect for flash operations)"}

        if not files:
            return {"success": False, "error": "No files specified"}

        # Validate all files exist and build address/path pairs
        flash_args: list[str] = []
        validated_files: list[dict] = []
        total_size = 0

        for entry in files:
            if "address" not in entry or "path" not in entry:
                return {
                    "success": False,
                    "error": f"Each file entry must have 'address' and 'path' keys. Got: {entry}",
                }

            fw_path = Path(entry["path"])
            if not fw_path.exists():
                return {"success": False, "error": f"File not found: {entry['path']}"}

            file_size = fw_path.stat().st_size
            total_size += file_size
            validated_files.append({
                "address": entry["address"],
                "path": str(fw_path),
                "size": file_size,
            })
            flash_args.extend([entry["address"], str(fw_path)])

        start_time = time.time()

        # Build esptool command: write-flash [options] addr1 file1 addr2 file2 ...
        # Options come after the subcommand in esptool v5.x
        args = ["write-flash"]
        if compress:
            args.append("--compress")
        if not verify:
            args.append("--no-verify")
        args.extend(flash_args)

        result = await self._run_esptool(port, args, timeout=300.0)

        if not result["success"]:
            return {
                "success": False,
                "error": result["error"],
                "port": port,
                "files": validated_files,
            }

        output = result["output"]
        elapsed = round(time.time() - start_time, 1)

        # Parse bytes written from output
        bytes_written = 0
        write_matches = re.findall(r"Wrote (\d+) bytes", output)
        for match in write_matches:
            bytes_written += int(match)

        verified = "Hash of data verified" in output or "Verified" in output

        return {
            "success": True,
            "port": port,
            "files": validated_files,
            "file_count": len(validated_files),
            "total_size": total_size,
            "bytes_written": bytes_written,
            "compressed": compress,
            "verified": verified if verify else None,
            "elapsed_seconds": elapsed,
        }

    async def _verify_flash_impl(
        self,
        context: Context,
        firmware_path: str,
        port: str | None,
        address: str,
    ) -> dict[str, Any]:
        """Verify flash contents match a file via esptool verify-flash."""

        fw_path = Path(firmware_path)
        if not fw_path.exists():
            return {"success": False, "error": f"File not found: {firmware_path}"}

        if not port:
            return {"success": False, "error": "Port is required (no auto-detect for flash operations)"}

        start_time = time.time()
        file_size = fw_path.stat().st_size

        # Use hyphenated command form (verify-flash not verify_flash)
        result = await self._run_esptool(
            port,
            ["verify-flash", address, str(fw_path)],
            timeout=120.0,
        )

        elapsed = round(time.time() - start_time, 1)

        # Check for verification success or failure in output
        output = result.get("output", "") or result.get("error", "")
        output_lower = output.lower()
        verified = (
            "verification successful" in output_lower
            or "verify ok" in output_lower
            or "digest matched" in output_lower
        )
        mismatch = (
            "verify failed" in output_lower
            or "mismatch" in output_lower
            or "does not match" in output_lower
        )

        if result["success"] and verified:
            return {
                "success": True,
                "verified": True,
                "port": port,
                "firmware_path": firmware_path,
                "address": address,
                "file_size": file_size,
                "elapsed_seconds": elapsed,
            }
        elif mismatch:
            # Extract mismatch details if available
            return {
                "success": True,
                "verified": False,
                "mismatch": True,
                "port": port,
                "firmware_path": firmware_path,
                "address": address,
                "file_size": file_size,
                "elapsed_seconds": elapsed,
                "details": output.strip() if output else None,
            }
        else:
            return {
                "success": False,
                "error": result.get("error", "Verification failed"),
                "port": port,
                "firmware_path": firmware_path,
                "address": address,
            }

    async def health_check(self) -> dict[str, Any]:
        """Component health check"""
        return {"status": "healthy", "note": "Flash manager ready"}
