from typing import Any, Self, Literal, overload
from collections.abc import Iterator
from abc import ABCMeta, abstractmethod
from dataclasses import dataclass
import re
import os
import sys
import hashlib
import pathlib
import shutil
import subprocess
import importlib.metadata
import importlib.resources
import sysconfig

from glasgow.support import logging
from glasgow.support.lazy import lazy


__all__ = ["ToolchainNotFound", "ToolOutOfDate", "Toolchain", "find_toolchain"]


logger = logging.getLogger(__name__)


class ToolchainNotFound(Exception):
    pass


class ToolOutOfDate(Exception):
    pass


class Tool(metaclass=ABCMeta):
    def __init__(self, name):
        self.name = str(name)

    @property
    def package_name(self) -> str:
        if self.name == "yosys" or self.name.startswith("nextpnr-"):
            return self.name
        if self.name == "icepack":
            return "nextpnr-ice40"
        if self.name == "ecppack":
            return "nextpnr-ecp5"
        raise NotImplementedError(f"package name for tool {self.name} is not known")

    @property
    def env_var_name(self) -> str:
        """Name of environment variable used by Amaranth to configure tool location."""
        # Identical to amaranth._toolchain.tool_env_var.
        return self.name.upper().replace("-", "_").replace("+", "X")

    @property
    @abstractmethod
    def command(self) -> str | None:
        """Command name for invoking the tool.

        Full path to the executable that can be used to run the tool, or ``None`` if the tool
        is not available.
        """
        raise NotImplementedError

    @property
    @abstractmethod
    def available(self) -> bool:
        """Tool availability.

        ``True`` if the tool is installed, ``False`` otherwise. Installed binary may still not
        be runnable, or might be too old to be useful.
        """
        raise NotImplementedError

    @property
    @abstractmethod
    def version(self) -> tuple[str, ...] | None:
        """Tool version number.

        ``None`` if version number could not be determined, or a tool-specific tuple if it could.
        """
        raise NotImplementedError

    @property
    @abstractmethod
    def identifier(self) -> bytes | None:
        """Unique tool identifier.

        Returns an array of 16 bytes that uniquely identifies the behavior of this particular tool
        in its entirety, but has no other meaning. Typically implemented by hashing the binary and
        its data files.
        """

    def __repr__(self) -> str:
        return f"<{self.__class__.__module__}.{self.__class__.__name__} {self.name}>"


class WasmTool(Tool):
    PREFIX = "yowasp-"

    @property
    def python_package(self) -> str:
        return self.PREFIX + self.package_name

    @property
    def available(self) -> bool:
        try:
            importlib.metadata.metadata(self.python_package)
        except importlib.metadata.PackageNotFoundError:
            return False
        else:
            return True

    @property
    def command(self) -> str | None:
        if not self.available:
            return None

        basename = self.PREFIX + self.name
        # We cannot assume that the command is on PATH and accessible by its basename. This
        # will not be true when Glasgow is running from a pipx virtual environment (which isn't
        # activated when the `glasgow` script is run). Also, our build environment does not
        # even *have* PATH.
        match os.name:
            case "nt":
                schemes = ["nt_venv", "nt_user", "nt"]
            case "posix":
                schemes = ["posix_venv", "posix_user", "posix_home", "posix_prefix"]
        for scheme in schemes:
            script_path  = os.path.join(sysconfig.get_path("scripts", scheme), basename)
            script_path += sysconfig.get_config_var("EXE")
            if os.path.exists(script_path):
                return script_path
        else:
            raise FileNotFoundError(f"script {basename!r} not found; "
                                    f"this is an issue with your installation")

    @property
    def version(self) -> tuple[int, ...] | None:
        if not self.available:
            return None

        # Running Wasm tools for the first time can incur a significant delay, so use
        # the version from the Python package metadata (which is guaranteed to be the same).
        # This makes querying the version at least as fast as for the native tools.
        a, b, _c, *d = importlib.metadata.version(self.python_package).split(".")
        # This is a bit of a mess. In practice, the third component of the Wasm tool version
        # doesn't correspond to anything useful for versions we care about, so skip it to make
        # native and Wasm tools return more or less the same version. The last component is
        # the git commit distance, which is vital for properly handling pre-release packages.
        return (a, b, *d)

    @property
    def identifier(self) -> bytes | None:
        if not self.available:
            return None

        hasher = hashlib.blake2s()
        for file_entry in importlib.metadata.files(self.python_package):
            if file_entry.hash is None:
                continue # RECORD, *.pyc, etc
            hasher.update(file_entry.hash.value.encode("utf-8"))
        return hasher.digest()[:16]


class SystemTool(Tool):
    @staticmethod
    def get_output(args: list[str]) -> str:
        return subprocess.run(args,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            encoding="utf-8").stdout.strip()

    @property
    def available(self) -> bool:
        return self.command is not None

    @property
    def command(self) -> str | None:
        return shutil.which(os.environ.get(self.env_var_name, self.name))

    @property
    def version(self) -> tuple[str, ...] | None:
        if not self.available:
            return None

        if self.name == "yosys":
            # Yosys 0.26+50 (git sha1 ef8ed21a2, ccache clang 11.0.1-2 -O0 -fPIC)
            raw_version = self.get_output([self.command, "--version"])
            if matches := re.match(r"^Yosys ([^\s]+) \(git sha1 ([0-9a-f]+)", raw_version):
                version = matches[1].replace("+", ".").split(".")
                return (*version, "g" + matches[2])

        elif self.name.startswith("nextpnr-"):
            # nextpnr-ice40 -- Nex... ...ce and Route (Version nextpnr-0.2-48-g20e595e2)
            raw_version = self.get_output([self.command, "--version"])
            if matches := re.match(r".+?\(Version .+?-(.+)\)$", raw_version):
                return (*matches[1].replace("-", ".").split("."),)

        elif self.name == "icepack":
            # does not have versions; does not have an option to return version
            return ("0",)

        elif self.name == "ecppack":
            # Project Trellis ecppack Version 1.3-3-g6845f33
            raw_version = self.get_output([self.command, "--version"])
            if matches := re.match(r".+?Version (.+)$", raw_version):
                return (*matches[1].replace("-", ".").split("."),)

        else:
            raise NotImplementedError(f"Cannot extract version from tool {self.name!r}")

        # Unparseable version; presume unavailable.
        return None

    def _iter_data_files(self):
        if not self.available:
            return None

        def iter_files(start):
            for root, _dirs, files in os.walk(start):
                for file in files:
                    yield os.path.join(root, file)

        if self.name == "yosys":
            # In the odd case of `GLASGOW_TOOLCHAIN=system YOSYS=yowasp-yosys` (which does still
            # happen), we can't access importlib metadata. However, the version should never be
            # the same for two different builds, so we don't have to hash data files.
            if self.command.endswith("yowasp-yosys"):
                return iter([])
            # Check if we're dealing with an installed Yosys binary, where the share files
            # are co-located in the same prefix. Sometimes the share directory path will be
            # baked into Yosys and sometimes it will be referenced to /proc/self/exe, but
            # the result is the same.
            yosys_prefix = pathlib.Path(self.command).parent.parent
            yosys_datdir = yosys_prefix.joinpath("share", "yosys")
            if yosys_datdir.exists:
                return iter_files(yosys_datdir)
            # If we have `yosys-config`, this is probably a development version of Yosys. If not,
            # we don't have any way to hash the data directory, and continuing is unsafe.
            if yosys_datdir := self.get_output([f"{self.command}-config", "--datdir"]):
                return iter_files(yosys_datdir)
            else:
                raise FileNotFoundError("could not find Yosys data directory")
        else:
            # It is unclear if it is feasible to get at the data files and other dependencies
            # for nextpnr. However, while it is possible to ship chipdb separately (and Wasm
            # builds do so), the overwhelming majority of native builds is likely shipping
            # nxetpnr as a self-contained binary that loads no data.
            #
            # Icepack is a self-contained binary. Ecppack is similar to nextpnr.
            #
            # This is likely fine.
            return iter([])

    _identifier_cache: str | None = None

    @property
    def identifier(self) -> bytes | None:
        if not self.available:
            return None

        if self._identifier_cache is not None:
            return self._identifier_cache

        hasher = hashlib.blake2s()
        with open(self.command, "rb") as file:
            hasher.update(file.read())
        for data_filename in self._iter_data_files():
            with open(data_filename, "rb") as file:
                hasher.update(file.read())
        self._identifier_cache = hasher.digest()[:16]
        return self._identifier_cache


class JsTool(Tool):
    @property
    def _js_bridge(self) -> Any:
        import js
        return js.glasgowToolchain

    @property
    def available(self) -> bool:
        return self._js_bridge.available(self.package_name)

    @property
    def command(self) -> str | None:
        if self.available:
            return self.name
        else:
            return None

    @property
    def version(self) -> tuple[str, ...] | None:
        if not self.available:
            return None

        # Running Wasm tools can incur a significant delay; request the version from
        # the toolchain bridge object.
        return tuple(self._js_bridge.version(self.package_name).split("."))

    @property
    def identifier(self) -> bytes | None:
        if not self.available:
            return None

        # This implementation assumes that no two different tool builds will ever have the same
        # version metadata. This is less robust than hashing every tool component, but it's not
        # practical to do the latter on JS hosts.
        hasher = hashlib.blake2s()
        hasher.update(self.name.encode("utf-8"))
        hasher.update(self._js_bridge.version(self.package_name).encode("utf-8"))
        return hasher.digest()[:16]


@dataclass
class _NixPackage:
    name:    str
    version: tuple[str, ...]
    command: str

    @classmethod
    def parse(cls) -> dict[str, Self]:
        with importlib.resources.files(__package__).joinpath("toolchain_nix.txt").open() as f:
            packages = []
            for line in f.readlines():
                if m := re.match(r"^\s*(#.*)?$", line):
                    pass # whitespace or comment only
                elif m := re.match(r"^\s*(\S+)\s+(\S+)\s+(\S+)\s*(#.*)?$", line):
                    packages.append(cls(name=m[1], version=tuple(m[2].split(".")), command=m[3]))
                else:
                    assert False, f"invalid toolchain_nix.txt syntax: {line!r}"
            return {package.name: package for package in packages}


class NixTool(Tool):
    _STORE = _NixPackage.parse()

    @property
    def available(self) -> bool:
        return self.package_name in self._STORE

    @property
    def command(self) -> str | None:
        if package := self._STORE[self.package_name]:
            return package.command
        else:
            return None

    @property
    def version(self) -> tuple[str, ...] | None:
        if package := self._STORE[self.package_name]:
            return package.version
        else:
            return None

    @property
    def identifier(self) -> tuple[str, ...] | None:
        if package := self._STORE[self.package_name]:
            # The store path uniquely identifies a specific build of the tool.
            hasher = hashlib.blake2s()
            hasher.update(package.command.encode("utf-8"))
            return hasher.digest()[:16]
        else:
            return None


class Toolchain:
    def __init__(self, tools):
        self.tools = list(tools)

    @property
    def available(self) -> bool:
        """Toolchain availability.

        ``True`` if every tool is available, ``False`` otherwise.
        """
        return all(tool.available for tool in self.tools)

    @property
    def missing(self) -> Iterator[str]:
        """Tools that are missing from the toolchain.

        An iterator that yields the name of every tool whose version could not be determined,
        because it is either unavailable or crashes when run.
        """
        return (tool.name for tool in self.tools if not tool.available or tool.version is None)

    @property
    def env_vars(self) -> dict[str, str]:
        """Environment variables to bring the toolchain in scope.

        An environment dictionary that includes entries for every of the tools included in this
        toolchain, and nothing else.

        Can be passed to :meth:`amaranth.build.run.BuildPlan.execute_local()` as the `env`
        argument in order to build a bitstream using this toolchain while isolating the build
        from any environmental impurity.
        """
        return {tool.env_var_name: tool.command for tool in self.tools}

    @property
    def versions(self) -> dict[str, tuple[str, ...]]:
        """Versions of tools.

        A dictionary that maps names of tools to their versions.
        """
        return {tool.name: tool.version for tool in self.tools}

    @property
    def identifier(self) -> bytes | None:
        """Unique toolchain identifier.

        Returns an array of 16 bytes that uniquely identifies this particular collection of tools,
        but has no other meaning.
        """
        hasher = hashlib.blake2s()
        for tool in self.tools:
            if not tool.available:
                return None
            hasher.update(tool.identifier)
        return hasher.digest()[:16]

    def assert_version(self, tool: str, required: tuple[str, ...]):
        # This is pretty janky, but unfortunately we don't have a clean way to implement this.
        def comparable(version: tuple[str, ...]):
            return [int(chunk) for chunk in version[:len(required)]]
        assert comparable(("0", "9")) < comparable(("0", "91"))

        versions = self.versions
        if tool in versions and comparable(versions[tool]) < comparable(required):
            raise ToolOutOfDate(
                f"tool {tool!r} is out of date: installed version {'.'.join(versions[tool])} "
                f"is older than required version {'.'.join(required)}")

    def __str__(self) -> str:
        return ", ".join(f"{name} {'.'.join(ver or ('(unavailable)',))}"
                         for name, ver in self.versions.items())

    def __repr__(self) -> str:
        return (f"<{self.__class__.__module__}.{self.__class__.__name__} " +
                " ".join(f"{tool.command}=={'.'.join(tool.version or ('unavailable',))}"
                         for tool in self.tools) +
                f">")


_ALL_TOOLS = ["yosys", "nextpnr-ice40", "icepack", "nextpnr-ecp5", "ecppack"]


@overload
def find_toolchain(tools=_ALL_TOOLS, *, quiet: Literal[True]) -> Toolchain | None:
    pass


def find_toolchain(tools=_ALL_TOOLS, *, quiet: Literal[False] = False) -> Toolchain:
    """Discover a toolchain.

    Returns a :class:`Toolchain` that includes all of the requested tools chosen according to
    the ``GLASGOW_TOOLCHAIN`` environment variable, or raises :exn:`ToolchainNotFound` if such
    toolchain isn't available within the constraints.
    """
    env_var_name = "GLASGOW_TOOLCHAIN"
    available_toolchains = {}
    if sys.platform == "emscripten":
        available_toolchains["js"]      = Toolchain(map(JsTool,     tools))
    else:
        available_toolchains["nix"]     = Toolchain(map(NixTool,    tools))
        available_toolchains["builtin"] = Toolchain(map(WasmTool,   tools))
        available_toolchains["system"]  = Toolchain(map(SystemTool, tools))

    kinds = os.environ.get(env_var_name, ",".join(available_toolchains.keys())).split(",")
    for kind in kinds:
        if kind not in available_toolchains:
            if quiet:
                return None
            logger.error(f"the {env_var_name} environment variable contains "
                         f"an unrecognized toolchain kind {kind!r}, available: "
                         f"{', '.join(available_toolchains)}")
            raise ToolchainNotFound(f"Unknown toolchain kind {kind!r} in {env_var_name}")

    selected_toolchains = {kind: available_toolchains[kind] for kind in kinds}
    for kind, toolchain in selected_toolchains.items():
        if toolchain.available:
            logger.debug(f"using toolchain {kind!r} ({toolchain})")
            for tool in toolchain.tools:
                logger.trace(f"tool {tool.name!r} is invoked as {tool.command!r}")
            logger.trace(f"toolchain ID is %s",
                lazy(lambda toolchain=toolchain: toolchain.identifier.hex()))
            for tool in toolchain.tools:
                logger.trace(f"tool ID of {tool.name!r} is %s",
                    lazy(lambda tool=tool: tool.identifier.hex()))
            return toolchain

    else:
        if quiet:
            return None
        examined = ", ".join(f"{kind} (missing {', '.join(toolchain.missing)})"
                             for kind in kinds)
        if env_var_name in os.environ:
            logger.error(f"could not find a usable FPGA toolchain; "
                         f"examined (according to {env_var_name}): {examined}")
        else:
            logger.error(f"could not find a usable FPGA toolchain; examined: {examined}")
            logger.error(f"consider reinstalling the package with the 'builtin-toolchain' "
                         f"feature enabled, "
                         f"e.g.: `pipx install --force -e glasgow/software[builtin-toolchain]`")
        raise ToolchainNotFound(f"No usable toolchain is available (examined: {', '.join(kinds)})")
