import os
import re
import shlex
import shutil
import subprocess
import zipfile
from dataclasses import dataclass
from pathlib import Path

import yaml


class BuildError(RuntimeError):
    """Raised when a build stage fails or produces invalid artifacts."""


SYNTHESIS_STAGES = ("vitis", "vivado", "both")
FAT_BLOCKS = ("DATA_TYPES", "MODEL_CONFIG", "NUM_LAYERS")
BUILD_ROOT_DIRNAME = ".cache"


@dataclass(frozen=True)
class BuildPaths:
    fat_dir: Path
    repo_root: Path
    xilinx_root: Path
    hls_template_dir: Path
    vivado_template_dir: Path
    hls_script: Path
    vivado_script: Path
    vitis_settings: Path
    vivado_settings: Path
    hls_source: Path


@dataclass(frozen=True)
class StagedBuildPaths:
    build_root: Path
    hls_dir: Path
    vivado_dir: Path
    hls_script: Path
    vivado_script: Path
    vivado_project_dir: Path
    hls_source: Path
    hls_export_zip: Path
    hls_ip_repo_dir: Path
    vivado_out_dir: Path


def _env_path(name: str, default: Path) -> Path:
    value = os.environ.get(name)
    return Path(value).expanduser().resolve() if value else default.resolve()


def _first_existing_path(candidates: list[Path]) -> Path:
    for candidate in candidates:
        resolved = candidate.expanduser().resolve()
        if resolved.is_file():
            return resolved
    return candidates[0].expanduser().resolve()


def _settings_path(env_name: str, candidates: list[Path]) -> Path:
    value = os.environ.get(env_name)
    if value:
        return Path(value).expanduser().resolve()
    return _first_existing_path(candidates)


def _settings_missing_message(label: str, env_name: str, path: Path) -> str:
    return (
        f"{label} not found: {path}. "
        f"Set {env_name} to your Xilinx settings64.sh path, for example "
        f"{env_name}=/opt/Xilinx/{label.split()[0]}/2022.2/settings64.sh"
    )


def resolve_build_paths() -> BuildPaths:
    """Resolve vendor-tool and project-template paths from FAT_* environment variables."""
    fat_dir = Path(__file__).resolve().parent
    repo_root = fat_dir
    xilinx_root = _env_path("FAT_XILINX_ROOT", repo_root / "xilinx")
    hls_dir = _env_path("FAT_HLS_DIR", xilinx_root / "hls")
    vivado_dir = _env_path("FAT_VIVADO_DIR", xilinx_root / "vivado")

    vitis_settings_candidates = [
        repo_root / "tools" / "Vitis_HLS" / "2022.2" / "settings64.sh",
        repo_root.parent / "tools" / "Vitis_HLS" / "2022.2" / "settings64.sh",
        repo_root.parent.parent / "tools" / "Vitis_HLS" / "2022.2" / "settings64.sh",
        Path("/opt/Xilinx/Vitis_HLS/2022.2/settings64.sh"),
        Path("/tools/Xilinx/Vitis_HLS/2022.2/settings64.sh"),
        Path("/tools/Vitis_HLS/2022.2/settings64.sh"),
        Path.home() / "Xilinx" / "Vitis_HLS" / "2022.2" / "settings64.sh",
    ]
    vivado_settings_candidates = [
        repo_root / "tools" / "Vivado" / "2022.2" / "settings64.sh",
        repo_root.parent / "tools" / "Vivado" / "2022.2" / "settings64.sh",
        repo_root.parent.parent / "tools" / "Vivado" / "2022.2" / "settings64.sh",
        Path("/opt/Xilinx/Vivado/2022.2/settings64.sh"),
        Path("/tools/Xilinx/Vivado/2022.2/settings64.sh"),
        Path("/tools/Vivado/2022.2/settings64.sh"),
        Path.home() / "Xilinx" / "Vivado" / "2022.2" / "settings64.sh",
    ]

    return BuildPaths(
        fat_dir=fat_dir,
        repo_root=repo_root,
        xilinx_root=xilinx_root,
        hls_template_dir=hls_dir,
        vivado_template_dir=vivado_dir,
        hls_script=_env_path("FAT_HLS_SCRIPT", hls_dir / "run.sh"),
        vivado_script=_env_path("FAT_VIVADO_TCL", vivado_dir / "run.tcl"),
        vitis_settings=_settings_path("FAT_VITIS_SETTINGS", vitis_settings_candidates),
        vivado_settings=_settings_path("FAT_VIVADO_SETTINGS", vivado_settings_candidates),
        hls_source=_env_path("FAT_HLS_SOURCE", hls_dir / "src" / "siso.cpp"),
    )


def _require_file(path: Path, label: str) -> None:
    if not path.is_file():
        raise BuildError(f"{label} not found: {path}")


def _require_dir(path: Path, label: str) -> None:
    if not path.is_dir():
        raise BuildError(f"{label} not found: {path}")


def _quote(path: Path) -> str:
    return shlex.quote(str(path))


def _run_stage(stage_name: str, shell_command: str, cwd: Path) -> None:
    print(f"[build:{stage_name}] Starting")
    process = subprocess.Popen(
        ["bash", "-lc", shell_command],
        cwd=cwd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
    )

    assert process.stdout is not None
    for line in process.stdout:
        print(f"[build:{stage_name}] {line}", end="")

    return_code = process.wait()
    if return_code != 0:
        raise BuildError(f"{stage_name} failed with exit code {return_code}")

    print(f"[build:{stage_name}] Completed")


def _require_artifact(path: Path, label: str) -> Path:
    if not path.is_file():
        raise BuildError(f"No {label} artifact found at {path}")
    return path


def _resolve_hwh_artifact(out_dir: Path, bitstream: Path) -> Path:
    hwh_path = out_dir / "encoder.hwh"
    if hwh_path.is_file():
        return hwh_path

    xsa_path = out_dir / "encoder.xsa"
    if not xsa_path.is_file():
        raise BuildError(f"No HWH artifact found at {hwh_path} and no XSA fallback found at {xsa_path}")

    try:
        with zipfile.ZipFile(xsa_path) as archive:
            hwh_members = [name for name in archive.namelist() if name.endswith(".hwh")]
            if not hwh_members:
                raise BuildError(f"XSA does not contain an .hwh file: {xsa_path}")
            bitstem = bitstream.stem
            matching_members = [
                name for name in hwh_members if Path(name).stem == bitstem
            ]

            if len(matching_members) == 1:
                member = matching_members[0]
            elif len(hwh_members) == 1:
                member = hwh_members[0]
            elif matching_members:
                raise BuildError(
                    f"XSA contains multiple .hwh files matching bitstream '{bitstem}', cannot choose automatically: {xsa_path}"
                )
            else:
                raise BuildError(
                    f"XSA contains multiple .hwh files and none match bitstream '{bitstem}': {xsa_path}"
                )
            hwh_path.write_bytes(archive.read(member))
    except zipfile.BadZipFile as exc:
        raise BuildError(f"Failed to read XSA as a zip archive: {xsa_path}") from exc

    print(f"Extracted {member} from {xsa_path} -> {hwh_path}")
    return hwh_path


def _copy_artifact(source: Path, destination: Path) -> None:
    if destination.exists():
        destination.unlink()
    shutil.copy2(source, destination)
    print(f"Copied {source} -> {destination}")


def _replace_fat_block(contents: str, block_name: str, replacement: str) -> str:
    pattern = re.compile(
        rf"(?P<start>// FAT_BEGIN {block_name}\n)(?P<body>.*?)(?P<end>// FAT_END {block_name})",
        re.DOTALL,
    )
    match = pattern.search(contents)
    if match is None:
        raise BuildError(f"Missing FAT marker block '{block_name}' in HLS source")

    replacement_body = replacement.rstrip() + "\n"
    return contents[:match.start("body")] + replacement_body + contents[match.end("body"):]


def _load_model_yaml(data_dir: Path) -> dict:
    config_path = data_dir / "model.yaml"
    if not config_path.is_file():
        raise BuildError(f"Could not read model.yaml from {config_path}")
    with open(config_path, "r") as f:
        return yaml.safe_load(f) or {}


def _render_hls_blocks(config: dict) -> dict[str, str]:
    kernel = config.get("kernel", {})
    model = config.get("model", {})

    store_type = "half" if kernel.get("store_fp16", False) else "float"
    acc_type = "half" if kernel.get("acc_fp16", False) else "float"
    sequence_length = int(model["sequence_length"])
    d_model = int(model["d_model"])
    num_heads = int(model["num_heads"])
    ff_dim = int(model["ff_dim"])
    num_layers = int(kernel.get("encoder_layer_count", 1))

    return {
        "DATA_TYPES": (
            f"typedef {store_type} data_t;\n"
            f"typedef {acc_type} acc_t;"
        ),
        "MODEL_CONFIG": (
            f"constexpr int SEQ_LEN = {sequence_length};\n"
            f"constexpr int D_MODEL = {d_model};\n"
            "constexpr int QKV_DIM = 3 * D_MODEL;\n"
            f"constexpr int H = {num_heads};\n"
            "constexpr int d_k = D_MODEL / H;\n"
            "constexpr acc_t LN_EPSILON = (acc_t)1e-5f;\n"
            f"constexpr int FF_DIM = {ff_dim};"
        ),
        "NUM_LAYERS": (
            "#ifndef NUM_LAYERS\n"
            f"#define NUM_LAYERS {num_layers}\n"
            "#endif"
        ),
    }


def _stage_tree(source_dir: Path, destination_dir: Path) -> None:
    if destination_dir.exists():
        shutil.rmtree(destination_dir)
    shutil.copytree(source_dir, destination_dir)


def _stage_build_workspace(target_dir: Path, paths: BuildPaths, stage: str) -> StagedBuildPaths:
    build_root = Path(os.environ.get("FAT_BUILD_ROOT", paths.repo_root / BUILD_ROOT_DIRNAME)).expanduser().resolve()
    experiment_name = target_dir.name
    stage_root = build_root / experiment_name
    hls_stage_dir = stage_root / "hls"
    vivado_stage_dir = stage_root / "vivado"

    stage_root.mkdir(parents=True, exist_ok=True)

    if stage in {"vitis", "both"} or not hls_stage_dir.exists():
        _stage_tree(paths.hls_template_dir, hls_stage_dir)
    _stage_tree(paths.vivado_template_dir, vivado_stage_dir)

    return StagedBuildPaths(
        build_root=stage_root,
        hls_dir=hls_stage_dir,
        vivado_dir=vivado_stage_dir,
        hls_script=hls_stage_dir / paths.hls_script.relative_to(paths.hls_template_dir),
        vivado_script=vivado_stage_dir / paths.vivado_script.relative_to(paths.vivado_template_dir),
        vivado_project_dir=vivado_stage_dir / "project",
        hls_source=hls_stage_dir / paths.hls_source.relative_to(paths.hls_template_dir),
        hls_export_zip=hls_stage_dir / "export" / "export.zip",
        hls_ip_repo_dir=hls_stage_dir / "export" / "export",
        vivado_out_dir=Path(os.environ.get("FAT_VIVADO_OUT_DIR", stage_root / "out")).expanduser().resolve(),
    )


def _update_hls_source_from_model_yaml(data_dir: Path, hls_source: Path) -> None:
    """Render model.yaml values into the marked FAT blocks in the staged HLS source."""
    config = _load_model_yaml(data_dir)
    source_text = hls_source.read_text()
    for block_name, block_text in _render_hls_blocks(config).items():
        source_text = _replace_fat_block(source_text, block_name, block_text)
    hls_source.write_text(source_text)
    print(f"Updated HLS source parameters in {hls_source}")


def synthesize_fpga_kernel(
    data_dir: str,
    stage: str = "both",
    jobs: int | None = None,
    dry_run: bool = False,
) -> tuple[Path | None, Path | None]:
    target_dir = Path(data_dir).expanduser().resolve()
    if not target_dir.is_dir():
        raise BuildError(f"Data directory not found: {target_dir}")
    if stage not in SYNTHESIS_STAGES:
        raise BuildError(f"Unsupported synthesis stage '{stage}'. Expected one of: {', '.join(SYNTHESIS_STAGES)}")

    paths = resolve_build_paths()
    _require_file(paths.hls_script, "HLS wrapper")
    _require_file(paths.vivado_script, "Vivado Tcl script")
    _require_file(paths.hls_source, "HLS source file")

    staged = _stage_build_workspace(target_dir, paths, stage)
    _require_file(staged.hls_script, "staged HLS wrapper")
    _require_file(staged.vivado_script, "staged Vivado Tcl script")
    _require_file(staged.hls_source, "staged HLS source file")

    job_count = jobs if jobs is not None else int(os.environ.get("FAT_VIVADO_JOBS", "8"))

    hls_command = None
    vivado_command = None
    dry_run_warnings = []

    if stage in {"vitis", "both"}:
        if not paths.vitis_settings.is_file():
            message = _settings_missing_message("Vitis_HLS settings script", "FAT_VITIS_SETTINGS", paths.vitis_settings)
            if dry_run:
                dry_run_warnings.append(message)
            else:
                raise BuildError(message)
        hls_command = (
            f"source {_quote(paths.vitis_settings)} && "
            f"FAT_VITIS_SETTINGS={_quote(paths.vitis_settings)} "
            f"FAT_HLS_TCL_SCRIPT={_quote(staged.hls_dir / 'run_hls.tcl')} "
            f"FAT_HLS_EXPORT_ZIP={_quote(staged.hls_export_zip)} "
            f"{_quote(staged.hls_script)}"
        )

    if stage in {"vivado", "both"}:
        if not paths.vivado_settings.is_file():
            message = _settings_missing_message("Vivado settings script", "FAT_VIVADO_SETTINGS", paths.vivado_settings)
            if dry_run:
                dry_run_warnings.append(message)
            else:
                raise BuildError(message)
        if stage == "vivado" and not staged.hls_ip_repo_dir.is_dir():
            message = (
                "Vivado-only synthesis requires an existing staged HLS IP repository. "
                "Run '--synthesis', '--synthesis vitis', or provide a preserved "
                f"HLS workspace under {staged.hls_dir} first."
            )
            if dry_run:
                dry_run_warnings.append(message)
            else:
                raise BuildError(message)
        vivado_command = (
            f"source {_quote(paths.vivado_settings)} && "
            f"vivado -mode batch -source {_quote(staged.vivado_script)} "
            f"-tclargs {_quote(staged.hls_ip_repo_dir)} "
            f"{_quote(staged.vivado_project_dir)} {job_count} {_quote(staged.vivado_out_dir)}"
        )

    if dry_run:
        config = _load_model_yaml(target_dir)
        rendered_blocks = _render_hls_blocks(config)
        print("FAT FPGA synthesis dry run")
        print(f"Stage: {stage}")
        print(f"Data directory: {target_dir}")
        print(f"Xilinx asset root: {paths.xilinx_root}")
        print(f"Staged build root: {staged.build_root}")
        print(f"Staged HLS directory: {staged.hls_dir}")
        print(f"Staged Vivado directory: {staged.vivado_dir}")
        print(f"Staged Vivado project directory: {staged.vivado_project_dir}")
        print(f"HLS source template: {paths.hls_source}")
        print(f"Staged HLS source: {staged.hls_source}")
        print(f"Expected HLS export archive: {staged.hls_export_zip}")
        print(f"Expected HLS IP repository: {staged.hls_ip_repo_dir}")
        for block_name in FAT_BLOCKS:
            print(f"Planned {block_name}:\n{rendered_blocks[block_name]}")
        for warning in dry_run_warnings:
            print(f"WARNING: {warning}")
        if hls_command is not None:
            print(f"HLS command: {hls_command}")
        if vivado_command is not None:
            print(f"Vivado command: {vivado_command}")
        print(f"Expected Vivado output directory: {staged.vivado_out_dir}")
        return target_dir / "encoder.bit", target_dir / "encoder.hwh"

    if hls_command is not None:
        _update_hls_source_from_model_yaml(target_dir, staged.hls_source)
        _run_stage("hls", hls_command, staged.hls_dir)
        _require_file(staged.hls_export_zip, "HLS export archive")
        _require_dir(staged.hls_ip_repo_dir, "HLS IP repository")

    if vivado_command is None:
        return None, None

    _run_stage("vivado", vivado_command, staged.vivado_dir)
    bitstream = _require_artifact(staged.vivado_out_dir / "encoder.bit", "bitstream")
    hwh = _resolve_hwh_artifact(staged.vivado_out_dir, bitstream)

    target_bit = target_dir / "encoder.bit"
    target_hwh = target_dir / "encoder.hwh"
    _copy_artifact(bitstream, target_bit)
    _copy_artifact(hwh, target_hwh)

    return target_bit, target_hwh
