"""Flash programming module for Zynq Flasher GUI.

Handles wiping, programming, and verifying QSPI flash memory using
program_flash (part of Xilinx Vitis/SDK), which connects through
hw_server to the Zynq JTAG interface.

Compatible with Xilinx 2018.3 — 2023.2.
"""

from __future__ import annotations

import binascii
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

from config_manager import Config
from vitis_checker import get_program_flash_path

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


@dataclass
class FlashResult:
    """Result of a flash programming operation."""

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


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

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


def _run_command(
    cmd: list[str],
    timeout: int = 300,
    callback: ProgressCallback = None,
) -> subprocess.CompletedProcess[str]:
    """Run a subprocess command with real-time streaming output.

    Uses subprocess_utils.run_streaming to stream stdout/stderr in
    real time, parsing for warnings, errors, and progress.

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

    Returns:
        CompletedProcess result.
    """
    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


def _build_program_flash_cmd_base(
    flash_tool: str,
    config: Config,
) -> list[str]:
    """Build the base command for program_flash.

    program_flash requires an FSBL ELF to initialize the flash controller
    on Zynq devices.

    Args:
        flash_tool: Path to program_flash executable.
        config: Application configuration.

    Returns:
        Base command list without the operation-specific flags.
    """
    return [
        flash_tool,
        "-f", config.bootloader_bin_path,
        "-fsbl", config.fsbl_elf_path,
        "-flash_type", config.flash_type,
    ]


# ── Public API ──────────────────────────────────────────────────────


def wipe_flash(
    config: Config,
    callback: ProgressCallback = None,
) -> FlashResult:
    """Erase QSPI flash before programming.

    - erase_all=False: -erase_only (matches BIN size, works on all chips)
    - erase_all=True:  creates a temp file of full flash capacity, uses
      -erase_only with it. Avoids -erase_all flag which fails on
      s25fl256s1 (no memory density info).

    Args:
        config: Application configuration (needs fsbl_elf_path).
        callback: Optional progress callback.

    Returns:
        FlashResult with status.
    """
    flash_tool = get_program_flash_path(
        xilinx_root=getattr(config, "xilinx_path", ""),
    )
    if not flash_tool:
        return FlashResult(
            step="wipe",
            success=False,
            message="program_flash not found in PATH",
        )

    fsbl_path = Path(config.fsbl_elf_path) if config.fsbl_elf_path else None
    if not fsbl_path or not fsbl_path.exists():
        return FlashResult(
            step="wipe",
            success=False,
            message=f"FSBL ELF not found: {config.fsbl_elf_path or '(not configured)'}",
        )

    if callback:
        callback("start", "Erasing QSPI flash...")

    erase_all = getattr(config, 'erase_all', False)

    try:
        if erase_all:
            flash_size = _flash_capacity(config)
            import tempfile, os
            tmp = tempfile.NamedTemporaryFile(delete=False, suffix='.bin')
            try:
                tmp.seek(flash_size - 1)
                tmp.write(b'\x00')
                tmp.close()
                cmd = [
                    flash_tool,
                    "-f", tmp.name,
                    "-fsbl", str(fsbl_path),
                    "-flash_type", config.flash_type,
                    "-erase_only",
                ]
                result = _run_command(
                    cmd,
                    timeout=config.step_timeouts.get("flash_erase", 120),
                    callback=callback,
                )
            finally:
                os.unlink(tmp.name)
        else:
            cmd = _build_program_flash_cmd_base(flash_tool, config)
            cmd += ["-erase_only"]
            result = _run_command(
                cmd,
                timeout=config.step_timeouts.get("flash_erase", 120),
                callback=callback,
            )

        success = result.returncode == 0
        return FlashResult(
            step="wipe",
            success=success,
            message="Flash erased successfully" if success else "Flash erase failed",
            output=(result.stdout + result.stderr)[:4000],
        )
    except subprocess.TimeoutExpired:
        return FlashResult(step="wipe", success=False, message="Flash erase timed out")
    except FileNotFoundError:
        return FlashResult(step="wipe", success=False, message="program_flash not found")


def _flash_capacity(config: Config) -> int:
    """Get flash chip capacity in bytes from flash_model.

    Args:
        config: Application configuration.

    Returns:
        Capacity in bytes (default 32 MiB).
    """
    model = getattr(config, 'flash_model', 's25fl256s1').lower()
    capacities: dict[str, int] = {
        's25fl256': 32 * 1024 * 1024,
        's25fl128': 16 * 1024 * 1024,
        's25fl064': 8 * 1024 * 1024,
    }
    for key, cap in capacities.items():
        if key in model:
            return cap
    return 32 * 1024 * 1024


def program_flash_bin(
    config: Config,
    bin_path: str | Path,
    callback: ProgressCallback = None,
) -> FlashResult:
    """Program the QSPI flash with a bootloader binary.

    Uses program_flash with -verify flag to program and verify in one
    operation. Does NOT erase first (call wipe_flash before this).

    Args:
        config: Application configuration.
        bin_path: Path to the .bin bootloader image.
        callback: Optional progress callback.

    Returns:
        FlashResult with status.
    """
    bin_path = Path(bin_path)
    if not bin_path.exists():
        return FlashResult(
            step="program",
            success=False,
            message=f"BIN file not found: {bin_path}",
        )

    flash_tool = get_program_flash_path(
        xilinx_root=getattr(config, "xilinx_path", ""),
    )
    if not flash_tool:
        return FlashResult(
            step="program",
            success=False,
            message="program_flash not found in PATH",
        )

    if callback:
        callback("start", f"Programming flash: {bin_path.name} ({bin_path.stat().st_size / 1024 / 1024:.1f} MB)...")

    # Override the BIN path in the command (use provided bin_path)
    cmd = [
        flash_tool,
        "-f", str(bin_path),
        "-fsbl", config.fsbl_elf_path,
        "-flash_type", config.flash_type,
        "-verify",
        "-no_erase",
    ]

    try:
        result = _run_command(cmd, timeout=config.step_timeouts.get("flash_program", 600), callback=callback)
        success = result.returncode == 0
        return FlashResult(
            step="program",
            success=success,
            message="Flash programmed successfully" if success else "Flash programming failed",
            output=(result.stdout + result.stderr)[:4000],
        )
    except subprocess.TimeoutExpired:
        return FlashResult(step="program", success=False, message="Flash programming timed out (600s)")


def verify_flash(
    config: Config,
    bin_path: str | Path,
    callback: ProgressCallback = None,
) -> FlashResult:
    """Verify the programmed flash content against a local BIN file.

    Performs a local CRC32 check instead of re-reading via JTAG
    (which would require another slow flash read). For a full hardware
    verify, run program_flash with -verify -no_program.

    Args:
        config: Application configuration (unused, kept for API consistency).
        bin_path: Path to the original .bin file.
        callback: Optional progress callback.

    Returns:
        FlashResult with CRC32 status.
    """
    bin_path = Path(bin_path)
    if not bin_path.exists():
        return FlashResult(
            step="verify",
            success=False,
            message=f"File not found: {bin_path}",
        )

    if callback:
        callback("start", "Computing local CRC32 for verification...")

    try:
        original_crc = _compute_crc32(bin_path)
        if callback:
            callback("progress", f"CRC32: {original_crc:#010x}")
        return FlashResult(
            step="verify",
            success=True,
            message=f"Local CRC32: {original_crc:#010x}",
        )
    except OSError as e:
        return FlashResult(
            step="verify",
            success=False,
            message=f"Read error: {e}",
        )


def _compute_crc32(path: Path) -> int:
    """Compute CRC32 checksum of a file.

    Reads the file in chunks to handle large binaries.

    Args:
        path: Path to the file.

    Returns:
        CRC32 value as unsigned 32-bit integer.
    """
    crc = 0
    with open(path, "rb") as f:
        while chunk := f.read(65536):
            crc = binascii.crc32(chunk, crc)
    return crc & 0xFFFFFFFF


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


def full_flash_program(
    config: Config,
    bin_path: str | Path,
    callback: ProgressCallback = None,
) -> list[FlashResult]:
    """Execute full flash workflow: wipe → program → verify.

    Erase scope is controlled by config.erase_all:
    - True: entire flash chip (uses -erase_all)
    - False: sectors matching BIN file size (uses -erase_only)

    Args:
        config: Application configuration.
        bin_path: Path to the .bin bootloader image.
        callback: Optional progress callback.

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

    # Step 1: Erase (scope determined by config.erase_all)
    results.append(wipe_flash(config, callback))
    if not results[-1].success:
        return results

    # Step 2: Program and verify (program_flash -verify)
    results.append(program_flash_bin(config, bin_path, callback))
    if not results[-1].success:
        return results

    # Step 3: Local CRC check
    results.append(verify_flash(config, bin_path, callback))
    return results