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

Uses Vivado batch-mode TCL to scan the JTAG chain via hw_server,
returning device names, IDCODEs, and PS/PL presence.
"""

from __future__ import annotations

import re
import socket
import subprocess
import tempfile
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable

from config_manager import Config
from vitis_checker import get_vivado_path, get_hw_server_path, build_tool_command


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

    name: str
    idcode: str = ""       # IDCODE in hex (e.g. "03736093")
    is_zynq: bool = False
    is_arm_dap: bool = False
    is_pl_device: bool = False


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

    name: str
    part_number: str       # e.g. "XC7Z100"
    family: str            # e.g. "zynq7"
    speed: str             # e.g. "-1"
    idcode: str = ""
    is_programmable: bool = False


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

    devices: list[ZynqDevice]
    jtag_chain: list[JtagDevice]
    is_present: bool
    ps_detected: bool
    pl_detected: bool
    hw_server_url: str
    error: str = ""


# ── Patterns ──────────────────────────────────────────────────────────
# Vivado TCL output:  DEVICE:arm_dap_0 IDCODE:0100101110100...0111
_DEVICE_LINE_RE = re.compile(r"^DEVICE:(\S+)\s+IDCODE:([01]+)")

# Zynq PL: xc7z100, xc7z100_0, xc7z100_1, etc.
_ZYNQ_RE = re.compile(r"^xc\d+z\d+[a-z]?\d+", re.IGNORECASE)

# ARM DAP (PS): arm_dap, arm_dap_0, arm_dap_1
_ARM_DAP_RE = re.compile(r"^arm_dap", re.IGNORECASE)

HW_SERVER_PORT = 3121


# ── Helpers ───────────────────────────────────────────────────────────

def _port_open(host: str = "localhost", port: int = HW_SERVER_PORT) -> bool:
    """Check if a TCP port is accepting connections."""
    try:
        with socket.create_connection((host, port), timeout=1):
            return True
    except OSError:
        return False


def _idcode_to_hex(binary_str: str) -> str:
    """Convert Vivado's binary IDCODE to uppercase hex."""
    try:
        return f"{int(binary_str, 2):08X}"
    except ValueError:
        return binary_str


# ── Main entry ────────────────────────────────────────────────────────

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

    Lifecycle:
      1. If hw_server is already running (port 3121 open) → reuse it
      2. If not running → start it, kill it when done
      3. Run Vivado batch TCL to enumerate JTAG devices

    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}"
    xilinx_root = getattr(config, "xilinx_path", "")

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

    # ── 1. Find tools ──────────────────────────────────────────
    vivado_path = get_vivado_path(xilinx_root=xilinx_root)
    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",
        )

    hw_server_path = get_hw_server_path(xilinx_root=xilinx_root)
    hw_server_was_running = _port_open()

    # ── 2. Start hw_server if needed ───────────────────────────
    hw_proc: subprocess.Popen | None = None
    if not hw_server_was_running:
        if not hw_server_path:
            return ZynqCheckResult(
                devices=[], jtag_chain=[], is_present=False,
                ps_detected=False, pl_detected=False,
                hw_server_url=hw_server_url,
                error="hw_server not found — install Vivado/Vitis",
            )
        try:
            hw_proc = subprocess.Popen(
                build_tool_command(hw_server_path),
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
            # Wait for port to become available
            for _ in range(30):  # 3s max
                if _port_open():
                    break
                time.sleep(0.1)
            else:
                return ZynqCheckResult(
                    devices=[], jtag_chain=[], is_present=False,
                    ps_detected=False, pl_detected=False,
                    hw_server_url=hw_server_url,
                    error="hw_server started but port 3121 not open",
                )
        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"Failed to start hw_server: {e}",
            )

    # ── 3. Run Vivado TCL ──────────────────────────────────────
    try:
        tcl = _build_jtag_query_tcl()
        if callback:
            callback("progress", "Scanning JTAG chain...")

        python_timeout = config.step_timeouts.get("check_env", 120)
        result = subprocess.run(
            build_tool_command(vivado_path, "-mode", "batch", "-source", tcl),
            capture_output=True, text=True, timeout=python_timeout,
        )
        # Clean up temp TCL file
        try:
            Path(tcl).unlink(missing_ok=True)
        except Exception:
            pass

        output = result.stdout + result.stderr
        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")

        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=f"JTAG scan timed out after {config.step_timeouts.get('check_env', 120)}s",
        )
    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:
        # ── 4. Kill hw_server only if WE started it ────────────
        if hw_proc is not None and hw_proc.poll() is None:
            try:
                hw_proc.terminate()
                hw_proc.wait(timeout=3)
            except subprocess.TimeoutExpired:
                hw_proc.kill()
                hw_proc.wait()
            except OSError:
                pass


# ── TCL builder ────────────────────────────────────────────────────────

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

    Returns:
        Path to the temporary TCL script file.
    """
    content = """\
open_hw_manager
connect_hw_server
open_hw_target
puts "===JTAG_START==="
foreach dev [get_hw_devices] {
  set name [get_property NAME $dev]
  set idcode [get_property IDCODE $dev]
  puts "DEVICE:$name IDCODE:$idcode"
}
puts "===JTAG_END==="
quit
"""
    import os
    path = Path(tempfile.gettempdir()) / f"zynq_jtag_{os.getpid()}.tcl"
    path.write_text(content)
    return str(path)


# ── Parsers ────────────────────────────────────────────────────────────

def _parse_jtag_chain(output: str) -> list[JtagDevice]:
    """Parse Vivado TCL output into JtagDevice list.

    Expected format:
        DEVICE:arm_dap_0 IDCODE:01001011101000000000010001110111
        DEVICE:xc7z100_1 IDCODE:00000011011100110110000010010011
    """
    devices: list[JtagDevice] = []
    section = _extract_section(output, "===JTAG_START===", "===JTAG_END===")
    if not section:
        return devices

    for line in section.splitlines():
        line = line.strip()
        if not line:
            continue
        match = _DEVICE_LINE_RE.match(line)
        if match:
            name = match.group(1)
            idcode_hex = _idcode_to_hex(match.group(2))
            devices.append(_classify_jtag_device(name, idcode_hex))

    return devices


def _filter_zynq_devices(jtag_chain: list[JtagDevice]) -> list[ZynqDevice]:
    """Filter Zynq devices from JTAG chain, propagating IDCODE."""
    zynq_devices: list[ZynqDevice] = []
    for jtag_dev in jtag_chain:
        if jtag_dev.is_zynq:
            zynq_dev = _parse_zynq_device(jtag_dev.name, jtag_dev.idcode)
            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 last occurrence of two markers."""
    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, idcode: str = "") -> JtagDevice:
    """Classify a JTAG device by name (Vivado format: xc7z100_1, arm_dap_0)."""
    is_zynq = bool(_ZYNQ_RE.match(device_name))
    is_arm_dap = bool(_ARM_DAP_RE.match(device_name))
    is_pl_device = is_zynq and not is_arm_dap

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


def _parse_zynq_device(device_name: str, idcode: str = "") -> ZynqDevice | None:
    """Parse Zynq device name (Vivado format: xc7z100_1)."""
    match = _ZYNQ_RE.match(device_name)
    if not match:
        return None

    # Strip index suffix: xc7z100_1 → XC7Z100
    base = device_name.rsplit("_", 1)[0] if "_" in device_name else device_name
    part_number = base.upper()

    family = "zynq7"
    speed = "-1"
    speed_match = re.search(r"(\d+)$", part_number)
    if speed_match:
        speed_num = int(speed_match.group(1))
        speed = "-2" if speed_num >= 2 else "-1"

    return ZynqDevice(
        name=device_name,
        part_number=part_number,
        family=family,
        speed=speed,
        idcode=idcode,
        is_programmable=True,
    )


# ── Status dict ────────────────────────────────────────────────────────

def get_zynq_status_dict(result: ZynqCheckResult) -> dict:
    """Convert ZynqCheckResult to a serializable dict."""
    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,
                "idcode": dev.idcode,
                "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
        ],
    }