"""
Production Tools Component

Provides factory programming, batch operations, quality control,
and production line integration tools.
"""

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 ProductionTools:
    """ESP production and factory programming tools"""

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

        @self.app.tool("esp_factory_program")
        async def factory_program(
            context: Context, program_config: dict[str, Any], port: str | None = None
        ) -> dict[str, Any]:
            """Program device for factory deployment"""
            return await self._factory_program_impl(context, program_config, port)

        @self.app.tool("esp_batch_program")
        async def batch_program(
            context: Context, device_list: list[str], firmware_path: str
        ) -> dict[str, Any]:
            """Program multiple devices in batch"""
            return await self._batch_program_impl(context, device_list, firmware_path)

        @self.app.tool("esp_quality_control")
        async def quality_control(
            context: Context, port: str | None = None, test_suite: str = "basic"
        ) -> dict[str, Any]:
            """Run quality control tests"""
            return await self._quality_control_impl(context, port, test_suite)

    async def _factory_program_impl(
        self,
        context: Context,
        program_config: dict[str, Any],
        port: str | None,
    ) -> dict[str, Any]:
        """Factory-program a device: erase → flash → verify.

        program_config should contain:
        {
            "firmware_path": "/path/to/firmware.bin",
            "address": "0x0",          # optional, default "0x0"
            "erase_before": true,       # optional, default true
            "verify": true,             # optional, default true
            "partition_table": "/path/to/partitions.bin",  # optional
            "partition_table_address": "0x8000",            # optional
            "bootloader": "/path/to/bootloader.bin",        # optional
            "bootloader_address": "0x1000",                 # optional
        }
        """

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

        firmware_path = program_config.get("firmware_path")
        if not firmware_path:
            return {"success": False, "error": "program_config must include 'firmware_path'"}

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

        erase_before = program_config.get("erase_before", True)
        verify = program_config.get("verify", True)
        address = program_config.get("address", "0x0")

        steps: list[dict[str, Any]] = []
        t_start = time.time()

        # Step 1: Erase flash
        if erase_before:
            result = await self._run_esptool(port, ["erase-flash"], timeout=60.0)
            steps.append({
                "step": "erase_flash",
                "success": result["success"],
                "error": result.get("error"),
            })
            if not result["success"]:
                return {
                    "success": False,
                    "error": f"Erase failed: {result['error']}",
                    "steps": steps,
                    "port": port,
                }

        # Step 2: Flash bootloader (if provided)
        bootloader = program_config.get("bootloader")
        if bootloader:
            bl_path = Path(bootloader)
            if not bl_path.exists():
                return {"success": False, "error": f"Bootloader not found: {bootloader}"}

            bl_addr = program_config.get("bootloader_address", "0x1000")
            result = await self._run_esptool(
                port,
                ["write-flash", bl_addr, bootloader],
                timeout=120.0,
            )
            steps.append({
                "step": "flash_bootloader",
                "address": bl_addr,
                "success": result["success"],
                "error": result.get("error"),
            })
            if not result["success"]:
                return {
                    "success": False,
                    "error": f"Bootloader flash failed: {result['error']}",
                    "steps": steps,
                    "port": port,
                }

        # Step 3: Flash partition table (if provided)
        partition_table = program_config.get("partition_table")
        if partition_table:
            pt_path = Path(partition_table)
            if not pt_path.exists():
                return {"success": False, "error": f"Partition table not found: {partition_table}"}

            pt_addr = program_config.get("partition_table_address", "0x8000")
            result = await self._run_esptool(
                port,
                ["write-flash", pt_addr, partition_table],
                timeout=120.0,
            )
            steps.append({
                "step": "flash_partition_table",
                "address": pt_addr,
                "success": result["success"],
                "error": result.get("error"),
            })
            if not result["success"]:
                return {
                    "success": False,
                    "error": f"Partition table flash failed: {result['error']}",
                    "steps": steps,
                    "port": port,
                }

        # Step 4: Flash main firmware
        write_args = ["write-flash"]
        if verify:
            write_args.append("--verify")
        write_args.extend([address, firmware_path])

        result = await self._run_esptool(port, write_args, timeout=300.0)
        steps.append({
            "step": "flash_firmware",
            "address": address,
            "success": result["success"],
            "error": result.get("error"),
        })
        if not result["success"]:
            return {
                "success": False,
                "error": f"Firmware flash failed: {result['error']}",
                "steps": steps,
                "port": port,
            }

        elapsed = round(time.time() - t_start, 2)

        return {
            "success": True,
            "port": port,
            "steps": steps,
            "total_time_seconds": elapsed,
            "firmware_path": firmware_path,
            "firmware_size_bytes": fw.stat().st_size,
        }

    async def _batch_program_impl(
        self,
        context: Context,
        device_list: list[str],
        firmware_path: str,
    ) -> dict[str, Any]:
        """Program multiple devices in parallel.

        Each device gets the same firmware flashed at 0x0 with erase + verify.
        Devices are programmed concurrently using asyncio.gather.
        """

        if not device_list:
            return {"success": False, "error": "device_list is empty"}

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

        t_start = time.time()

        async def program_one(port: str) -> dict[str, Any]:
            """Program a single device."""
            config = {
                "firmware_path": firmware_path,
                "erase_before": True,
                "verify": True,
            }
            return await self._factory_program_impl(context, config, port)

        # Run all programming tasks concurrently
        results = await asyncio.gather(
            *[program_one(port) for port in device_list],
            return_exceptions=True,
        )

        device_results = []
        succeeded = 0
        for port, result in zip(device_list, results, strict=True):
            if isinstance(result, Exception):
                device_results.append({
                    "port": port,
                    "success": False,
                    "error": str(result),
                })
            else:
                device_results.append({
                    "port": port,
                    "success": result.get("success", False),
                    "error": result.get("error"),
                    "time_seconds": result.get("total_time_seconds"),
                })
                if result.get("success"):
                    succeeded += 1

        elapsed = round(time.time() - t_start, 2)

        return {
            "success": succeeded == len(device_list),
            "total_devices": len(device_list),
            "succeeded": succeeded,
            "failed": len(device_list) - succeeded,
            "total_time_seconds": elapsed,
            "firmware_path": firmware_path,
            "devices": device_results,
        }

    async def _quality_control_impl(
        self,
        context: Context,
        port: str | None,
        test_suite: str,
    ) -> dict[str, Any]:
        """Run quality control checks on a device.

        Test suites:
        - "basic": chip-id, flash-id, read-mac (fast verification)
        - "extended": basic + flash read/verify + memory dump check
        """

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

        tests: list[dict[str, Any]] = []
        t_start = time.time()

        # Test 1: Chip identification
        result = await self._run_esptool(port, ["chip-id"], timeout=15.0)
        chip_info: dict[str, Any] = {"test": "chip_identification", "success": result["success"]}
        if result["success"]:
            output = 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)
            chip_info["chip"] = chip_match.group(1) if chip_match else "unknown"
            chip_info["chip_id"] = id_match.group(1) if id_match else "unknown"
        else:
            chip_info["error"] = result.get("error", "")[:200]
        tests.append(chip_info)

        # Test 2: Flash identification
        result = await self._run_esptool(port, ["flash-id"], timeout=15.0)
        flash_info: dict[str, Any] = {"test": "flash_identification", "success": result["success"]}
        if result["success"]:
            output = result["output"]
            mfr_match = re.search(r"Manufacturer:\s*(0x[0-9a-fA-F]+)", output)
            size_match = re.search(r"Detected flash size:\s*(\S+)", output)
            flash_info["manufacturer"] = mfr_match.group(1) if mfr_match else "unknown"
            flash_info["flash_size"] = size_match.group(1) if size_match else "unknown"
        else:
            flash_info["error"] = result.get("error", "")[:200]
        tests.append(flash_info)

        # Test 3: MAC address
        result = await self._run_esptool(port, ["read-mac"], timeout=15.0)
        mac_info: dict[str, Any] = {"test": "mac_address", "success": result["success"]}
        if result["success"]:
            mac_match = re.search(r"MAC:\s*([0-9a-fA-F:]+)", result["output"])
            mac_info["mac"] = mac_match.group(1) if mac_match else "unknown"
        else:
            mac_info["error"] = result.get("error", "")[:200]
        tests.append(mac_info)

        # Extended tests
        if test_suite == "extended":
            # Test 4: Read first 4KB of flash (checks flash connectivity)
            import tempfile

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

            try:
                result = await self._run_esptool(
                    port,
                    ["read-flash", "0x0", "4096", tmp_path],
                    timeout=60.0,
                )
                read_info: dict[str, Any] = {"test": "flash_read_4kb", "success": result["success"]}
                if result["success"]:
                    data = Path(tmp_path).read_bytes()
                    read_info["bytes_read"] = len(data)
                    # Check if flash is all 0xFF (erased) or has data
                    non_ff = sum(1 for b in data if b != 0xFF)
                    read_info["has_data"] = non_ff > 0
                    read_info["non_erased_bytes"] = non_ff
                else:
                    read_info["error"] = result.get("error", "")[:200]
                tests.append(read_info)
            finally:
                import os

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

        elapsed = round(time.time() - t_start, 2)

        # Determine overall pass/fail
        passed = sum(1 for t in tests if t["success"])
        all_passed = passed == len(tests)

        return {
            "success": True,
            "port": port,
            "test_suite": test_suite,
            "verdict": "PASS" if all_passed else "FAIL",
            "tests_run": len(tests),
            "tests_passed": passed,
            "tests_failed": len(tests) - passed,
            "total_time_seconds": elapsed,
            "tests": tests,
        }

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