"""
Diagnostics Component

Provides ESP device diagnostics including memory dumps, flash identification,
performance profiling, and comprehensive diagnostic reporting. All operations
shell out to esptool as async subprocesses.
"""

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

from fastmcp import Context, FastMCP

from ..config import ESPToolServerConfig

logger = logging.getLogger(__name__)

# Size suffixes for human-friendly parsing
_SIZE_MULTIPLIERS = {"B": 1, "KB": 1024, "MB": 1024 * 1024}


def _parse_size(size_str: str) -> int:
    """Parse a human-friendly size string like '1KB', '256B', '4MB' into bytes."""
    size_str = size_str.strip().upper()
    for suffix, mult in sorted(_SIZE_MULTIPLIERS.items(), key=lambda x: -len(x[0])):
        if size_str.endswith(suffix):
            num = size_str[: -len(suffix)].strip()
            return int(num) * mult
    # Try as plain integer (decimal or hex)
    return int(size_str, 0)


class Diagnostics:
    """ESP device diagnostics and analysis"""

    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 = 30.0,
    ) -> dict[str, Any]:
        """Run esptool as an async subprocess."""
        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()[:500]}

            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 diagnostic tools"""

        @self.app.tool("esp_memory_dump")
        async def memory_dump(
            context: Context,
            port: str | None = None,
            start_address: str = "0x0",
            size: str = "1KB",
        ) -> dict[str, Any]:
            """Dump device memory for analysis.

            Reads raw bytes from an arbitrary memory address on the ESP device
            using esptool's dump-mem command. Useful for inspecting bootloader
            state, peripheral registers, or RAM contents.

            The output is hex-formatted for readability. For flash memory reads,
            use esp_flash_read instead (faster, supports larger ranges).

            Args:
                port: Serial port or socket:// URI (required)
                start_address: Memory address to start reading (hex string, default: "0x0")
                size: Number of bytes to read (e.g. "256B", "1KB", "4KB", default: "1KB")
            """
            return await self._memory_dump_impl(context, port, start_address, size)

        @self.app.tool("esp_performance_profile")
        async def performance_profile(
            context: Context,
            port: str | None = None,
            duration: int = 30,
        ) -> dict[str, Any]:
            """Profile device communication performance.

            Measures serial transport speed by timing a sequence of esptool
            operations (chip-id, flash-id, small memory reads). Reports
            round-trip latencies and throughput estimates. Useful for comparing
            physical serial vs QEMU socket performance.

            Args:
                port: Serial port or socket:// URI (required)
                duration: Not used for timing control; kept for API compatibility
            """
            return await self._performance_profile_impl(context, port)

        @self.app.tool("esp_diagnostic_report")
        async def diagnostic_report(
            context: Context,
            port: str | None = None,
            include_memory: bool = False,
        ) -> dict[str, Any]:
            """Generate comprehensive diagnostic report for an ESP device.

            Collects chip identity, MAC address, flash information, and
            optionally a small memory dump into a single structured report.
            Useful for troubleshooting connectivity issues or characterizing
            an unknown device.

            For security-focused analysis, use esp_security_audit instead.

            Args:
                port: Serial port or socket:// URI (required)
                include_memory: Include a 256-byte memory dump from 0x0 (default: false)
            """
            return await self._diagnostic_report_impl(context, port, include_memory)

    async def _memory_dump_impl(
        self,
        context: Context,
        port: str | None,
        start_address: str,
        size: str,
    ) -> dict[str, Any]:
        """Read memory via esptool dump-mem (writes to temp file, then reads it)."""

        if not port:
            return {"success": False, "error": "Port is required for memory dump"}

        byte_count = _parse_size(size)
        if byte_count > 1024 * 1024:
            return {"success": False, "error": "Maximum dump size is 1MB"}

        # dump-mem writes raw bytes to a file
        import tempfile
        with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as tmp:
            tmp_path = tmp.name

        try:
            result = await self._run_esptool(
                port,
                ["dump-mem", start_address, str(byte_count), tmp_path],
                timeout=60.0,
            )

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

            # Read the dump file and format as hex
            from pathlib import Path
            dump_path = Path(tmp_path)
            if not dump_path.exists() or dump_path.stat().st_size == 0:
                return {"success": False, "error": "Dump file is empty", "port": port}

            raw = dump_path.read_bytes()

            # Format as hex dump (16 bytes per line with ASCII)
            hex_lines = []
            for offset in range(0, len(raw), 16):
                chunk = raw[offset : offset + 16]
                addr = int(start_address, 0) + offset
                hex_part = " ".join(f"{b:02x}" for b in chunk)
                ascii_part = "".join(chr(b) if 32 <= b < 127 else "." for b in chunk)
                hex_lines.append(f"0x{addr:08x}: {hex_part:<48s} {ascii_part}")

            return {
                "success": True,
                "port": port,
                "start_address": start_address,
                "bytes_read": len(raw),
                "hex_dump": "\n".join(hex_lines[:64]),  # Cap at 64 lines (1KB)
                "truncated": len(hex_lines) > 64,
            }

        finally:
            import os
            try:
                os.unlink(tmp_path)
            except OSError:
                pass

    async def _performance_profile_impl(
        self,
        context: Context,
        port: str | None,
    ) -> dict[str, Any]:
        """Profile serial transport by timing esptool operations."""

        if not port:
            return {"success": False, "error": "Port is required for profiling"}

        measurements: list[dict[str, Any]] = []

        # Test 1: chip-id (lightweight command)
        t0 = time.time()
        r = await self._run_esptool(port, ["chip-id"], timeout=15.0)
        elapsed = round(time.time() - t0, 3)
        measurements.append({
            "operation": "chip-id",
            "elapsed_seconds": elapsed,
            "success": r["success"],
        })

        # Test 2: flash-id (reads SPI flash ID register)
        t0 = time.time()
        r = await self._run_esptool(port, ["flash-id"], timeout=15.0)
        elapsed = round(time.time() - t0, 3)
        measurements.append({
            "operation": "flash-id",
            "elapsed_seconds": elapsed,
            "success": r["success"],
        })

        # Test 3: read-mac
        t0 = time.time()
        r = await self._run_esptool(port, ["read-mac"], timeout=15.0)
        elapsed = round(time.time() - t0, 3)
        measurements.append({
            "operation": "read-mac",
            "elapsed_seconds": elapsed,
            "success": r["success"],
        })

        # Test 4: read 4KB of flash (throughput test)
        import tempfile
        with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as tmp:
            tmp_path = tmp.name

        try:
            t0 = time.time()
            r = await self._run_esptool(
                port,
                ["read-flash", "0x0", "4096", tmp_path],
                timeout=60.0,
            )
            elapsed = round(time.time() - t0, 3)

            throughput = None
            if r["success"] and elapsed > 0:
                throughput = round(4096 / elapsed, 0)

            measurements.append({
                "operation": "read-flash (4KB)",
                "elapsed_seconds": elapsed,
                "success": r["success"],
                "throughput_bytes_per_sec": throughput,
            })
        finally:
            import os
            try:
                os.unlink(tmp_path)
            except OSError:
                pass

        # Summary
        successes = [m for m in measurements if m["success"]]
        avg_latency = (
            round(sum(m["elapsed_seconds"] for m in successes) / len(successes), 3)
            if successes
            else None
        )

        return {
            "success": True,
            "port": port,
            "measurements": measurements,
            "summary": {
                "operations_tested": len(measurements),
                "operations_succeeded": len(successes),
                "average_latency_seconds": avg_latency,
            },
        }

    async def _diagnostic_report_impl(
        self,
        context: Context,
        port: str | None,
        include_memory: bool,
    ) -> dict[str, Any]:
        """Generate comprehensive device diagnostic report."""

        if not port:
            return {"success": False, "error": "Port is required for diagnostic report"}

        report: dict[str, Any] = {"port": port}

        # 1. Chip identification
        chip_result = await self._run_esptool(port, ["chip-id"], timeout=15.0)
        if chip_result["success"]:
            output = chip_result["output"]
            chip_match = re.search(r"Chip is (\S+)", output)
            id_match = re.search(r"Chip ID:\s*(0x[0-9a-fA-F]+)", output)
            report["chip"] = chip_match.group(1) if chip_match else "unknown"
            report["chip_id"] = id_match.group(1) if id_match else "unknown"
        else:
            return {"success": False, "error": f"Cannot reach device: {chip_result['error']}", "port": port}

        # 2. MAC address
        mac_result = await self._run_esptool(port, ["read-mac"], timeout=15.0)
        if mac_result["success"]:
            mac_match = re.search(r"MAC:\s*([0-9a-fA-F:]+)", mac_result["output"])
            report["mac_address"] = mac_match.group(1) if mac_match else "unknown"

        # 3. Flash info
        flash_result = await self._run_esptool(port, ["flash-id"], timeout=15.0)
        if flash_result["success"]:
            output = flash_result["output"]
            mfr_match = re.search(r"Manufacturer:\s*(0x[0-9a-fA-F]+)", output)
            dev_match = re.search(r"Device:\s*(0x[0-9a-fA-F]+)", output)
            size_match = re.search(r"Detected flash size:\s*(\S+)", output)
            report["flash"] = {
                "manufacturer": mfr_match.group(1) if mfr_match else "unknown",
                "device": dev_match.group(1) if dev_match else "unknown",
                "size": size_match.group(1) if size_match else "unknown",
            }

        # 4. Optional memory dump
        if include_memory:
            mem_result = await self._memory_dump_impl(context, port, "0x0", "256B")
            if mem_result.get("success"):
                report["memory_dump_0x0"] = mem_result.get("hex_dump", "")

        report["success"] = True
        return report

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