"""Zynq JTAG presence checker for Zynq Flasher GUI.

Connects to hw_server via Vivado and scans the JTAG chain to detect
Zynq devices, returning chip model and presence information.
Checks for both PS (Processing System) and PL (Programmable Logic)
recognition on the JTAG chain.
"""

from __future__ import annotations

import os
import re
import shutil
import subprocess
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable

from config_manager import Config


@dataclass
class JtagDevice:
    """Information about a device on the JTAG chain."""

    name: str  # Device name from JTAG chain (e.g. "xc7z100_1")
    is_zynq: bool  # Whether this is a Zynq device
    is_arm_dap: bool  # Whether this is an ARM DAP (PS side)
    is_pl_device: bool  # Whether this is a PL device


@dataclass
class ZynqDevice:
    """Information about a detected Zynq device on the JTAG chain."""

    name: str  # Device name from JTAG chain (e.g. "xc7z100_1")
    part_number: str  # Extracted part number (e.g. "XC7Z100")
    family: str  # Family (e.g. "zynq7")
    speed: str  # Speed grade (e.g. "1", "2", "-1", "-2")
    idcode: str = ""  # JTAG IDCODE if available
    is_programmable: bool = False  # Whether the device is programmable


@dataclass
class ZynqCheckResult:
    """Result of Zynq device detection on the JTAG chain."""

    devices: list[ZynqDevice]  # All detected Zynq devices
    jtag_chain: list[JtagDevice]  # All devices on JTAG chain
    is_present: bool  # Whether any Zynq device was found
    ps_detected: bool  # Whether PS (ARM DAP) is detected
    pl_detected: bool  # Whether PL is detected
    hw_server_url: str  # hw_server connection URL used
    error: str = ""  # Error message if detection failed


# Zynq part number patterns from JTAG device names
# Examples: xc7z010_1, xc7z020_1, xc7z100_1, xc7z030_1
ZYNQ_PART_PATTERN = re.compile(
    r"^(xc\d+z\d+[a-z]?\d+)_\d+$", re.IGNORECASE
)

# ARM DAP pattern (PS side)
ARM_DAP_PATTERN = re.compile(r"^arm_dap_\d+$", re.IGNORECASE)

# hw_server default port
HW_SERVER_PORT = 3121


def check_zynq_jtag(config: Config, callback: Callable | None = None) -> ZynqCheckResult:
    """Check for Zynq devices on the JTAG chain via hw_server.

    Connects to hw_server using Vivado batch mode, scans the JTAG chain,
    and identifies Zynq devices by their JTAG device names. Also checks
    for PS (ARM DAP) and PL (Programmable Logic) recognition.

    Args:
        config: Application configuration.
        callback: Optional progress callback(status, message).

    Returns:
        ZynqCheckResult with detected devices and status.
    """
    hw_server_url = f"localhost:{HW_SERVER_PORT}"

    if callback:
        callback("start", "Connecting to hw_server...")

    # Find vivado executable
    vivado_path = _get_vivado_path(config.vitis_path)
    if not vivado_path:
        return ZynqCheckResult(
            devices=[],
            jtag_chain=[],
            is_present=False,
            ps_detected=False,
            pl_detected=False,
            hw_server_url=hw_server_url,
            error="Vivado not found",
        )

    # Create temporary TCL script
    tcl_content = _build_jtag_query_tcl()
    tcl_path = Path("/tmp/zynq_jtag_check_{}.tcl".format(os.getpid()))
    try:
        tcl_path.write_text(tcl_content)

        if callback:
            callback("progress", "Scanning JTAG chain...")

        # Run Vivado in batch mode
        cmd = [vivado_path, "-mode", "batch", "-source", str(tcl_path)]
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=30,
        )

        output = result.stdout + result.stderr

        # Parse JTAG device information from output
        jtag_chain = _parse_jtag_chain(output)
        devices = _filter_zynq_devices(jtag_chain)

        if callback:
            if devices:
                ps_note = "PS" if any(d.is_arm_dap for d in jtag_chain) else ""
                pl_note = "PL" if any(d.is_pl_device for d in jtag_chain) else ""
                notes = ", ".join(filter(None, [ps_note, pl_note]))
                callback("complete", f"Found {len(devices)} Zynq device(s) [{notes}]")
            else:
                callback("error", "No Zynq device found on JTAG chain")

        # Check PS and PL detection
        ps_detected = any(d.is_arm_dap for d in jtag_chain)
        pl_detected = any(d.is_pl_device for d in jtag_chain)

        return ZynqCheckResult(
            devices=devices,
            jtag_chain=jtag_chain,
            is_present=len(devices) > 0,
            ps_detected=ps_detected,
            pl_detected=pl_detected,
            hw_server_url=hw_server_url,
            error="" if devices else "No Zynq device found on JTAG chain",
        )

    except subprocess.TimeoutExpired:
        return ZynqCheckResult(
            devices=[],
            jtag_chain=[],
            is_present=False,
            ps_detected=False,
            pl_detected=False,
            hw_server_url=hw_server_url,
            error="JTAG scan timed out after 30s",
        )
    except OSError as e:
        return ZynqCheckResult(
            devices=[],
            jtag_chain=[],
            is_present=False,
            ps_detected=False,
            pl_detected=False,
            hw_server_url=hw_server_url,
            error=f"OS error: {e}",
        )
    finally:
        tcl_path.unlink(missing_ok=True)


def _build_jtag_query_tcl() -> str:
    """Build TCL script for JTAG device query.

    Returns:
        TCL script content as string.
    """
    return """
open_hw_manager
connect_hw_server
open_hw_target
puts "===JTAG_START==="
foreach dev [get_hw_devices] {
  set name [get_property NAME $dev]
  puts "DEVICE:$name"
}
puts "===JTAG_END==="
quit
"""


def _parse_jtag_chain(output: str) -> list[JtagDevice]:
    """Parse JTAG chain information from Vivado output.

    Args:
        output: Combined stdout/stderr from Vivado batch execution.

    Returns:
        List of JtagDevice objects for all devices on the JTAG chain.
    """
    devices: list[JtagDevice] = []

    # Extract device list from output
    devices_section = _extract_section(output, "===JTAG_START===", "===JTAG_END===")
    if not devices_section:
        return devices

    for line in devices_section.splitlines():
        line = line.strip()
        if line.startswith("DEVICE:"):
            device_name = line[len("DEVICE:"):]
            device = _classify_jtag_device(device_name)
            devices.append(device)

    return devices


def _filter_zynq_devices(jtag_chain: list[JtagDevice]) -> list[ZynqDevice]:
    """Filter Zynq devices from JTAG chain.

    Args:
        jtag_chain: List of all JTAG devices.

    Returns:
        List of ZynqDevice objects.
    """
    zynq_devices: list[ZynqDevice] = []
    for jtag_dev in jtag_chain:
        if jtag_dev.is_zynq:
            zynq_dev = _parse_zynq_device(jtag_dev.name)
            if zynq_dev:
                zynq_devices.append(zynq_dev)
    return zynq_devices


def _extract_section(
    text: str, start_marker: str, end_marker: str
) -> str | None:
    """Extract text between two markers, handling duplicate markers.

    Finds the last occurrence of each marker to avoid matching
    markers inside TCL script source code.

    Args:
        text: Full output text.
        start_marker: Starting marker string.
        end_marker: Ending marker string.

    Returns:
        Text between markers, or None if not found.
    """
    # Find last occurrence of each marker to avoid matching
    # markers inside TCL script source code
    start_idx = text.rfind(start_marker)
    end_idx = text.rfind(end_marker)
    if start_idx == -1 or end_idx == -1 or start_idx >= end_idx:
        return None
    return text[start_idx + len(start_marker):end_idx].strip()


def _classify_jtag_device(device_name: str) -> JtagDevice:
    """Classify a JTAG device by its name.

    Args:
        device_name: Device name from JTAG chain.

    Returns:
        JtagDevice with classification information.
    """
    is_zynq = bool(ZYNQ_PART_PATTERN.match(device_name))
    is_arm_dap = bool(ARM_DAP_PATTERN.match(device_name))
    is_pl_device = is_zynq and not is_arm_dap

    return JtagDevice(
        name=device_name,
        is_zynq=is_zynq,
        is_arm_dap=is_arm_dap,
        is_pl_device=is_pl_device,
    )


def _parse_zynq_device(device_name: str) -> ZynqDevice | None:
    """Parse a Zynq device from its JTAG chain name.

    Args:
        device_name: Device name from JTAG chain (e.g. "xc7z100_1").

    Returns:
        ZynqDevice if it's a Zynq chip, None otherwise.
    """
    match = ZYNQ_PART_PATTERN.match(device_name)
    if not match:
        return None

    part_base = match.group(1)
    # Convert to uppercase for part number
    part_number = part_base.upper()

    # Extract family and speed from part number
    # Examples: XC7Z010 -> family=zynq7, speed=1; XC7Z100 -> family=zynq7, speed=1
    family = "zynq7"
    speed = "1"  # Default speed grade

    # Parse speed grade from the numeric suffix
    speed_match = re.search(r"(\d+)$", part_base)
    if speed_match:
        speed_num = int(speed_match.group(1))
        # Map numeric speed to speed grade string
        if speed_num >= 2:
            speed = "-2"
        elif speed_num >= 1:
            speed = "-1"

    return ZynqDevice(
        name=device_name,
        part_number=part_number,
        family=family,
        speed=speed,
    )


def _get_vivado_path(vitis_path: str) -> str | None:
    """Find the Vivado executable path.

    Args:
        vitis_path: Vitis installation path.

    Returns:
        Path to Vivado executable, or None if not found.
    """
    path = shutil.which("vivado")
    if path:
        return path
    if vitis_path:
        candidate = Path(vitis_path) / "bin" / "vivado"
        if candidate.exists():
            return str(candidate)
        candidate = Path(vitis_path) / "bin" / "vivado.exe"
        if candidate.exists():
            return str(candidate)
    common_paths = [
        "/opt/xilinx/bin/vivado",
        os.path.expanduser("~/Xilinx/Vivado/bin/vivado"),
    ]
    for p in common_paths:
        if os.path.isfile(p):
            return p
    return None


def get_zynq_status_dict(result: ZynqCheckResult) -> dict:
    """Convert a ZynqCheckResult to a serializable dictionary.

    Args:
        result: ZynqCheckResult from check_zynq_jtag().

    Returns:
        Dictionary suitable for JSON serialization or GUI display.
    """
    return {
        "is_present": result.is_present,
        "ps_detected": result.ps_detected,
        "pl_detected": result.pl_detected,
        "hw_server_url": result.hw_server_url,
        "error": result.error,
        "jtag_chain": [
            {
                "name": dev.name,
                "is_zynq": dev.is_zynq,
                "is_arm_dap": dev.is_arm_dap,
                "is_pl_device": dev.is_pl_device,
            }
            for dev in result.jtag_chain
        ],
        "devices": [
            {
                "name": dev.name,
                "part_number": dev.part_number,
                "family": dev.family,
                "speed": dev.speed,
                "idcode": dev.idcode,
                "is_programmable": dev.is_programmable,
            }
            for dev in result.devices
        ],
    }