"""Configuration of HEIR compiler toolchain."""

import dataclasses
import os
import importlib.resources
import pathlib
from pathlib import Path
import shutil
from heir.backends.util.common import get_repo_root, is_pip_installed


dataclass = dataclasses.dataclass


@dataclass(frozen=True)
class HEIRConfig:
  heir_opt_path: str | Path
  heir_translate_path: str | Path
  techmap_dir_path: str | Path  # optional
  abc_path: str | Path  # optional


def development_heir_config() -> HEIRConfig:
  repo_root = get_repo_root()
  if not repo_root:
    return HEIRConfig(
        heir_opt_path="",
        heir_translate_path="",
        techmap_dir_path="",
        abc_path="",
    )

  techmap_dir_path = (
      repo_root
      / "bazel-bin"
      / "tools"
      / "heir-opt.runfiles"
      / "_main"
      / "lib"
      / "Transforms"
      / "YosysOptimizer"
      / "yosys"
  )
  if not techmap_dir_path.exists():
    techmap_dir_path = ""

  abc_path = (
      repo_root
      / "bazel-bin"
      / "tools"
      / "heir-opt.runfiles"
      / "abc"
      / "abc_bin"
  )
  if not abc_path.exists():
    abc_path = (
        repo_root
        / "bazel-bin"
        / "tools"
        / "heir-opt.runfiles"
        / "abc+"
        / "abc_bin"
    )
  if not abc_path.exists():
    abc_path = ""

  return HEIRConfig(
      heir_opt_path=repo_root / "bazel-bin" / "tools" / "heir-opt",
      heir_translate_path=repo_root / "bazel-bin" / "tools" / "heir-translate",
      techmap_dir_path=techmap_dir_path,
      abc_path=abc_path,
  )


def _resolve_binary(
    env_var_name: str,
    dev_path: Path | str,
    binary_name: str,
    required: bool = True,
) -> Path | str:
  # 1. Env Var
  val = os.environ.get(env_var_name)
  if val:
    return val

  # 2. Bazel-built dev path (if exists)
  dev_path_obj = Path(dev_path) if dev_path else None
  if dev_path_obj and dev_path_obj.exists():
    return dev_path_obj

  # 3. shutil.which PATH (if exists)
  which_path = shutil.which(binary_name)
  if which_path:
    return which_path

  # 4. dev path fallback (if it is defined and we want to fall back to it without error)
  if not required and dev_path:
    return dev_path

  # 5. RuntimeError
  if required:
    raise RuntimeError(
        f"Could not find binary '{binary_name}'. Please set {env_var_name} or"
        " ensure it is on your PATH or built via Bazel."
    )
  return ""


def from_os_env() -> HEIRConfig:
  """Create a HEIRConfig from environment variables.

  Note, this is required for running tests under bazel, as the locations
  of the various binaries are determined by bazel.

  The order of preference is:

  1. Environment variables HEIR_OPT_PATH, HEIR_TRANSLATE_PATH, HEIR_ABC_BINARY,
     or HEIR_YOSYS_SCRIPTS_DIR
  2. The default development configuration (relative to the project root, in
     bazel-bin) (if they exist)
  3. The path to the heir-opt or heir-translate binary on the PATH
  4. The dev path fallback (if not found elsewhere)

  Returns:
    The HEIRConfig
  """
  dev_config = development_heir_config()

  resolved_heir_opt_path = _resolve_binary(
      env_var_name="HEIR_OPT_PATH",
      dev_path=dev_config.heir_opt_path,
      binary_name="heir-opt",
      required=True,
  )
  resolved_heir_translate_path = _resolve_binary(
      env_var_name="HEIR_TRANSLATE_PATH",
      dev_path=dev_config.heir_translate_path,
      binary_name="heir-translate",
      required=True,
  )
  resolved_abc_path = _resolve_binary(
      env_var_name="HEIR_ABC_BINARY",
      dev_path=dev_config.abc_path,
      binary_name="abc",
      required=False,
  )

  resolved_techmap_dir_path = os.environ.get("HEIR_YOSYS_SCRIPTS_DIR")
  if not resolved_techmap_dir_path:
    if (
        dev_config.techmap_dir_path
        and Path(dev_config.techmap_dir_path).exists()
    ):
      resolved_techmap_dir_path = dev_config.techmap_dir_path
    else:
      resolved_techmap_dir_path = dev_config.techmap_dir_path

  return HEIRConfig(
      heir_opt_path=resolved_heir_opt_path,
      heir_translate_path=resolved_heir_translate_path,
      techmap_dir_path=resolved_techmap_dir_path,
      abc_path=resolved_abc_path,
  )


def from_pip_installation() -> HEIRConfig:
  """
  Configure HEIR binaries from the expected pip installation structure.
  """
  if not is_pip_installed():
    raise RuntimeError("HEIR is not installed via pip.")

  package_path = importlib.resources.files("heir")
  return HEIRConfig(
      heir_opt_path=package_path / "heir-opt",
      heir_translate_path=package_path / "heir-translate",
      # These paths are configured in setup.py
      techmap_dir_path=package_path / "techmaps",
      abc_path=package_path / "abc",
  )
