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

Starts hw_server explicitly, then uses Vivado batch-mode TCL to scan
the JTAG chain. On failure, Vivado's stdout/stderr is returned in the
error message for debugging.
"""

from __future__ import annotations

import re
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 = ""
    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
    family: str
    speed: str
    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 = ""
    hw_server_pid: int | None = None


# ── Patterns ──────────────────────────────────────────────────────────
_DEVICE_LINE_RE = re.compile(r"^DEVICE:(\S+)\s+IDCODE:([01]+)")
_ZYNQ_RE = re.compile(r"^xc\d+z\d+[a-z]?\d+", re.IGNORECASE)
_ARM_DAP_RE = re.compile(r"^arm_dap", re.IGNORECASE)
HW_SERVER_PORT = 3121


def _idcode_to_hex(binary_str: str) -> str:
    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.

    1. Start hw_server (Vivado may not auto-start on Windows)
    2. Run Vivado batch TCL to enumerate JTAG devices
    3. On failure, include Vivado output in error for debugging

    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 _fail("Vivado not found", hw_server_url)

    hw_server_path = get_hw_server_path(xilinx_root=xilinx_root)
    if not hw_server_path:
        return _fail("hw_server not found", hw_server_url)

    # ── 2. Start hw_server ─────────────────────────────────────
    hw_proc: subprocess.Popen | None = None
    hw_pid: int | None = None
    try:
        hw_proc = subprocess.Popen(
            build_tool_command(hw_server_path),
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        hw_pid = hw_proc.pid
        # Wait up to 10s for hw_server to bind
        started = False
        for _ in range(100):
            time.sleep(0.1)
            if _port_listening():
                started = True
                break
        if not started:
            return _fail("hw_server did not bind port 3121 within 10s", hw_server_url, pid=hw_pid)
    except OSError as e:
        return _fail(f"Failed to start hw_server: {e}", hw_server_url)

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

    # ── 3. Run Vivado TCL ──────────────────────────────────────
    tcl = _build_jtag_query_tcl()
    python_timeout = config.step_timeouts.get("check_env", 120)

    # Redirect Vivado journal/log to temp dir so we don't litter the cwd
    tmp = tempfile.gettempdir()
    vivado_args = [
        "-mode", "batch",
        "-tempDir", tmp,
        "-journal", str(Path(tmp) / "vivado_jtag.jou"),
        "-log", str(Path(tmp) / "vivado_jtag.log"),
        "-source", tcl,
    ]

    output_lines: list[str] = []
    try:
        cmd = build_tool_command(vivado_path, *vivado_args)
        # Use run() with communicate() under the hood — reliably
        # drains pipes on Windows even when .bat wrappers are involved
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=python_timeout,
        )
        output = result.stdout + result.stderr
        for line in output.splitlines():
            stripped = line.strip()
            if stripped:
                output_lines.append(stripped)
                if callback:
                    callback("progress", stripped)
    except subprocess.TimeoutExpired:
        return _fail(
            f"JTAG scan timed out after {python_timeout}s",
            hw_server_url,
            pid=hw_pid,
        )
    except OSError as e:
        return _fail(f"Failed to run Vivado: {e}", hw_server_url, pid=hw_pid)
    finally:
        try:
            Path(tcl).unlink(missing_ok=True)
        except Exception:
            pass
        for f in ("vivado_jtag.jou", "vivado_jtag.log"):
            try:
                (Path(tmp) / f).unlink(missing_ok=True)
            except Exception:
                pass

    output = "\n".join(output_lines)

    # Parse results
    jtag_chain = _parse_jtag_chain(output)
    devices = _filter_zynq_devices(jtag_chain)

    if devices:
        if callback:
            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:
        # Include Vivado output so user can debug
        err_detail = _extract_relevant(output)
        msg = "No Zynq device found on JTAG chain"
        if err_detail:
            msg += f" | Vivado: {err_detail[:300]}"
        if callback:
            callback("error", msg)
        return ZynqCheckResult(
            devices=[], jtag_chain=jtag_chain,
            is_present=False,
            ps_detected=False, pl_detected=False,
            hw_server_url=hw_server_url,
            error=msg,
            hw_server_pid=hw_pid,
        )

    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=True,
        ps_detected=ps_detected, pl_detected=pl_detected,
        hw_server_url=hw_server_url,
        error="",
        hw_server_pid=hw_pid,
    )


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

def _fail(error: str, url: str, pid: int | None = None) -> ZynqCheckResult:
    return ZynqCheckResult(
        devices=[], jtag_chain=[], is_present=False,
        ps_detected=False, pl_detected=False,
        hw_server_url=url, error=error,
        hw_server_pid=pid,
    )


def _port_listening(port: int = HW_SERVER_PORT) -> bool:
    """Check if hw_server is listening on its TCP port."""
    import socket
    try:
        with socket.create_connection(("localhost", port), timeout=1):
            return True
    except OSError:
        return False


def _extract_relevant(output: str) -> str:
    """Extract error/warning lines from Vivado output for user display."""
    lines = []
    for line in output.splitlines():
        line = line.strip()
        if not line:
            continue
        if any(kw in line for kw in ("ERROR", "WARNING", "CRITICAL", "FAIL", "cannot", "No")):
            lines.append(line)
    return " | ".join(lines[-5:])  # Last 5 relevant lines


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

def _build_jtag_query_tcl() -> str:
    """Build Vivado TCL script file for JTAG device enumeration."""
    content = """\
open_hw
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==="
close_hw
disconnect_hw_server
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."""
    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]:
    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:
    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:
    is_zynq = bool(_ZYNQ_RE.match(device_name))
    is_arm_dap = bool(_ARM_DAP_RE.match(device_name))
    return JtagDevice(
        name=device_name,
        idcode=idcode,
        is_zynq=is_zynq,
        is_arm_dap=is_arm_dap,
        is_pl_device=is_zynq and not is_arm_dap,
    )


def _parse_zynq_device(device_name: str, idcode: str = "") -> ZynqDevice | None:
    match = _ZYNQ_RE.match(device_name)
    if not match:
        return None
    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 = "-2" if int(speed_match.group(1)) >= 2 else "-1"
    return ZynqDevice(
        name=device_name, part_number=part_number,
        family=family, speed=speed,
        idcode=idcode, is_programmable=True,
    )


def get_zynq_status_dict(result: ZynqCheckResult) -> 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": d.name, "idcode": d.idcode,
             "is_zynq": d.is_zynq, "is_arm_dap": d.is_arm_dap,
             "is_pl_device": d.is_pl_device}
            for d in result.jtag_chain
        ],
        "devices": [
            {"name": d.name, "part_number": d.part_number,
             "family": d.family, "speed": d.speed,
             "idcode": d.idcode, "is_programmable": d.is_programmable}
            for d in result.devices
        ],
    }