"""Bitstream programming module for Zynq Flasher GUI.

Downloads a .bit bitstream to the Zynq PL fabric and runs a .elf
executable on the PS ARM core.

Uses xsct (Xilinx Software Command Line Tool) as the primary tool
for JTAG operations, with Vivado batch mode as a fallback.
Compatible with Xilinx 2018.3 — 2023.2.
"""

from __future__ import annotations

import os
import shutil
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

from config_manager import Config
from vitis_checker import get_xsct_path

# ── Types ───────────────────────────────────────────────────────────


@dataclass
class BitstreamResult:
    """Result of a bitstream programming or ELF execution operation."""

    step: str
    success: bool
    message: str
    output: str = ""


ProgressCallback = Callable[[str, str], None] | None

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


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

    Uses shutil.which() for PATH resolution, then common install paths.

    Args:
        vitis_path: Optional Vitis installation directory.

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


def _find_ps7_init_tcl(config: Config) -> str:
    """Find the ps7_init.tcl file for PS initialization.

    Checks:
    1. Same directory as bootloader BIT file (hw_platform)
    2. ps7_init_tcl_path in config (if added later)

    Args:
        config: Application configuration.

    Returns:
        Path to ps7_init.tcl, or empty string if not found.
    """
    # 1. Check config attribute (if added)
    if hasattr(config, 'ps7_init_tcl_path') and config.ps7_init_tcl_path:
        p = Path(config.ps7_init_tcl_path)
        if p.exists():
            return str(p)

    # 2. Same directory as BIT file (typical hw_platform layout)
    if config.bootloader_bit_path:
        bit_dir = Path(config.bootloader_bit_path).parent
        candidates = [
            bit_dir / "ps7_init.tcl",
        ]
        for c in candidates:
            if c.exists():
                return str(c)

    return ""


def _run_command(
    cmd: list[str],
    timeout: int = 120,
    callback: ProgressCallback = None,
) -> subprocess.CompletedProcess[str]:
    """Run a subprocess command with progress reporting.

    Args:
        cmd: Command and arguments.
        timeout: Maximum seconds.
        callback: Optional callback(status, message).

    Returns:
        CompletedProcess result.

    Raises:
        subprocess.TimeoutExpired: If command exceeds timeout.
    """
    from subprocess_utils import run_streaming

    if callback:
        callback("running", f"Executing: {' '.join(cmd)}")

    try:
        return run_streaming(cmd, timeout=timeout, callback=callback)
    except subprocess.TimeoutExpired:
        if callback:
            callback("error", f"{cmd[0]} timed out after {timeout}s")
        raise


# ── Bitstream Download ──────────────────────────────────────────────


def program_bitstream(
    config: Config,
    bit_path: str | Path,
    callback: ProgressCallback = None,
) -> BitstreamResult:
    """Download a .bit bitstream to the Zynq PL (FPGA fabric).

    Inits PS (ps7_init) before FPGA config to keep the JTAG DAP alive,
    allowing subsequent ELF download without reconnection issues.

    Args:
        config: Application configuration.
        bit_path: Path to the .bit file.
        callback: Optional progress callback.

    Returns:
        BitstreamResult with status.
    """
    bit_path = Path(bit_path)
    if not bit_path.exists():
        return BitstreamResult(
            step="bitstream",
            success=False,
            message=f"Bitstream file not found: {bit_path}",
        )

    if callback:
        callback("start", f"Loading bitstream: {bit_path.name}...")

    # Find ps7_init.tcl for PS initialization
    ps7_tcl = _find_ps7_init_tcl(config)

    # Try xsct first (preferred)
    xsct = get_xsct_path(
        vitis_path=config.vitis_path if hasattr(config, 'vitis_path') else "",
        xilinx_root=config.xilinx_path if hasattr(config, 'xilinx_path') else "",
    )
    if xsct:
        timeout = config.step_timeouts.get("bitstream", 60)
        return _program_with_xsct(xsct, bit_path, callback, ps7_tcl, timeout)

    # Fallback to Vivado
    vivado = _get_vivado_path(
        config.vitis_path if hasattr(config, 'vitis_path') else ""
    )
    if vivado:
        return _program_with_vivado(vivado, bit_path, callback)

    return BitstreamResult(
        step="bitstream",
        success=False,
        message="Neither xsct nor vivado found in PATH",
    )


def _program_with_xsct(
    xsct_path: str,
    bit_path: Path,
    callback: ProgressCallback = None,
    ps7_init_tcl: str = "",
    timeout: int = 60,
) -> BitstreamResult:
    """Program bitstream using xsct's fpga command.

    Inits the PS (clocks, DDR, MIO) before FPGA configuration to keep
    the JTAG DAP alive. Without PS init, the DAP breaks after fpga,
    preventing subsequent ELF download.

    Args:
        xsct_path: Path to xsct executable.
        bit_path: Path to .bit file.
        callback: Optional progress callback.
        ps7_init_tcl: Path to ps7_init.tcl for PS initialization.
        timeout: Subprocess timeout in seconds.

    Returns:
        BitstreamResult with status.
    """
    if callback:
        callback("progress", "Initializing PS then programming FPGA...")

    # Build TCL: PS init first, then FPGA config
    lines = ["connect"]
    if ps7_init_tcl and Path(ps7_init_tcl).exists():
        lines.append("targets 1")
        lines.append(f'source "{ps7_init_tcl}"')
        lines.append("ps7_init")
        if callback:
            callback("progress", "PS initialized (clocks, DDR, MIO)")
    lines.append(f'fpga -file "{bit_path}"')
    lines.append("exit")
    tcl_script = "\n".join(lines)

    import tempfile, os
    with tempfile.NamedTemporaryFile(mode='w', suffix='.tcl', delete=False) as f:
        f.write(tcl_script)
        tcl_file = f.name

    try:
        result = subprocess.run(
            [xsct_path, tcl_file],
            capture_output=True,
            text=True,
            timeout=timeout,
        )
        os.unlink(tcl_file)
        success = result.returncode == 0
        output = (result.stdout + result.stderr).strip()

        if callback:
            callback(
                "complete" if success else "error",
                "Bitstream " + ("loaded" if success else "failed"),
            )

        return BitstreamResult(
            step="bitstream",
            success=success,
            message="Bitstream loaded successfully" if success else "Bitstream load failed",
            output=output[:4000],
        )
    except subprocess.TimeoutExpired:
        os.unlink(tcl_file)
        return BitstreamResult(
            step="bitstream",
            success=False,
            message="xsct bitstream download timed out",
        )


def _program_with_vivado(
    vivado_path: str,
    bit_path: Path,
    callback: ProgressCallback = None,
) -> BitstreamResult:
    """Program bitstream using Vivado Hardware Manager.

    Args:
        vivado_path: Path to vivado executable.
        bit_path: Path to .bit file.
        callback: Optional progress callback.

    Returns:
        BitstreamResult with status.
    """
    if callback:
        callback("progress", "Using Vivado Hardware Manager...")

    tcl_content = f"""
open_hw_manager
connect_hw_server
open_hw_target
current_hw_device [lindex [get_hw_devices] 0]
program_hw_device -file "{bit_path}"
refresh_hw_device
quit
"""
    tcl_path = bit_path.parent / "temp_program.tcl"
    try:
        with open(tcl_path, "w") as f:
            f.write(tcl_content)

        result = _run_command(
            [vivado_path, "-mode", "batch", "-source", str(tcl_path)],
            timeout=120,
            callback=callback,
        )
        success = result.returncode == 0
        output = (result.stdout + result.stderr).strip()

        return BitstreamResult(
            step="bitstream",
            success=success,
            message="Bitstream loaded successfully" if success else "Bitstream load failed",
            output=output[:4000],
        )
    finally:
        tcl_path.unlink(missing_ok=True)


# ── ELF Execution ───────────────────────────────────────────────────


def run_elf(
    config: Config,
    elf_path: str | Path,
    callback: ProgressCallback = None,
) -> BitstreamResult:
    """Download and run a .elf executable on the Zynq PS ARM core.

    Uses xsct's `dow` (download) command, with Vivado as fallback.

    Args:
        config: Application configuration.
        elf_path: Path to the .elf file.
        callback: Optional progress callback.

    Returns:
        BitstreamResult with status.
    """
    elf_path = Path(elf_path)
    if not elf_path.exists():
        return BitstreamResult(
            step="elf",
            success=False,
            message=f"ELF file not found: {elf_path}",
        )

    if callback:
        callback("start", f"Running ELF: {elf_path.name}...")

    # Use xsct for CPU operations (Vivado HM cannot program PS CPU)
    xsct = get_xsct_path(
        vitis_path=config.vitis_path if hasattr(config, 'vitis_path') else "",
        xilinx_root=config.xilinx_path if hasattr(config, 'xilinx_path') else "",
    )
    if not xsct:
        return BitstreamResult(
            step="elf",
            success=False,
            message="xsct not found in PATH (required for CPU operations)",
        )

    # Auto-detect ps7_init.tcl from config or BIT file directory
    ps7_tcl = _find_ps7_init_tcl(config)
    timeout = config.step_timeouts.get("elf_download", 60)
    return _run_elf_with_xsct(xsct, elf_path, callback, ps7_tcl, timeout)


def _run_elf_with_xsct(
    xsct_path: str,
    elf_path: Path,
    callback: ProgressCallback = None,
    ps7_init_tcl: str = "",
    timeout: int = 60,
) -> BitstreamResult:
    """Download and run ELF using xsct's dow command.

    PS was already initialized by the prior BIT download step, so the
    DAP should be alive. First attempt skips ps7_init to avoid resetting
    PS state. If DAP is broken, falls back to re-init PS.

    Args:
        xsct_path: Path to xsct executable.
        elf_path: Path to .elf file.
        callback: Optional progress callback.
        ps7_init_tcl: Optional ps7_init.tcl path for recovery.
        timeout: Subprocess timeout in seconds.

    Returns:
        BitstreamResult with status.
    """
    import tempfile, os

    has_ps7 = bool(ps7_init_tcl and Path(ps7_init_tcl).exists())

    def _run(tcl: str) -> tuple[bool, str]:
        """Run TCL via xsct, return (success, output)."""
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".tcl", delete=False
        ) as tf:
            tf.write(tcl)
            tcl_file = tf.name
        try:
            result = subprocess.run(
                [xsct_path, tcl_file],
                capture_output=True, text=True, timeout=timeout,
            )
            os.unlink(tcl_file)
            out = (result.stdout + result.stderr).strip()
            ok = result.returncode == 0 or any(
                p in out for p in (
                    "Download successful", "Done", "done",
                    "Successfully downloaded",
                )
            )
            return ok, out
        except subprocess.TimeoutExpired:
            os.unlink(tcl_file)
            raise

    # ── Attempt 1: no ps7_init (DAP alive from BIT step) ──
    if callback:
        callback("progress", "Downloading ELF (DAP should be alive)...")

    tcl1 = "\n".join([
        "connect",
        "targets 2",
        "dow \"%s\"" % elf_path,
        "con",
        "exit",
    ])
    try:
        ok1, out1 = _run(tcl1)
    except subprocess.TimeoutExpired:
        return BitstreamResult(step="elf", success=False,
                               message="ELF download timed out")

    if ok1:
        if callback:
            callback("complete", "ELF executed")
        return BitstreamResult(step="elf", success=True,
                               message="ELF executed successfully",
                               output=out1[:4000])

    # ── Attempt 2: re-init PS to recover DAP ──
    if not has_ps7:
        if callback:
            callback("error", "ELF failed — no ps7_init.tcl for recovery")
        return BitstreamResult(step="elf", success=False,
                               message="ELF execution failed",
                               output=out1[:4000])

    if callback:
        callback(
            "progress",
            "DAP broken, re-initializing PS then retrying ELF...",
        )

    tcl2 = "\n".join([
        "connect",
        "targets 1",
        "source \"%s\"" % ps7_init_tcl,
        "ps7_init",
        "targets 2",
        "dow \"%s\"" % elf_path,
        "con",
        "exit",
    ])
    try:
        ok2, out2 = _run(tcl2)
    except subprocess.TimeoutExpired:
        return BitstreamResult(step="elf", success=False,
                               message="ELF download (recovery) timed out")

    if callback:
        callback("complete" if ok2 else "error",
                 "ELF " + ("executed (recovery)" if ok2 else "failed after recovery"))

    return BitstreamResult(
        step="elf",
        success=ok2,
        message=("ELF executed (DAP recovery)" if ok2 else "ELF execution failed"),
        output=out2[:4000],
    )


# ── Full Workflow ───────────────────────────────────────────────────


def full_bitstream_program(
    config: Config,
    bit_path: str | Path,
    elf_path: str | Path,
    callback: ProgressCallback = None,
) -> list[BitstreamResult]:
    """Execute full bootloader workflow: program bitstream → run ELF.

    Args:
        config: Application configuration.
        bit_path: Path to the .bit file.
        elf_path: Path to the .elf file.
        callback: Optional progress callback.

    Returns:
        List of BitstreamResult for each step.
    """
    results: list[BitstreamResult] = []

    results.append(program_bitstream(config, bit_path, callback))
    if not results[-1].success:
        return results

    results.append(run_elf(config, elf_path, callback))
    return results