"""
Common utilities for this package.

@date: 30.06.2025
@author: Baptiste Pestourie
"""

from __future__ import annotations

import importlib.util
import importlib
import keyword
import logging
import os
import shutil
import sys
import sysconfig
import tempfile
import warnings
from contextlib import contextmanager
from dataclasses import dataclass, field, fields
from enum import Enum, auto
from pathlib import Path
from types import ModuleType
from typing import (
    Generator,
    Literal,
    TypeGuard,
    NamedTuple,
    NewType,
    Self,
    assert_never,
    cast,
    overload,
)

from setuptools import Extension

from smelt.context import get_context

_logger = logging.getLogger(__name__)

PathExists = NewType("PathExists", Path)
ModuleName = NewType("ModuleName", str)
type _TomlData = str | list[_TomlData] | dict[str, _TomlData]
type TomlData = dict[str, _TomlData]


# --- Errors ---
class SmeltError(Exception):
    """
    Base exception class for Smelt errors
    """


class SmeltConfigError(SmeltError):
    """
    Raised when running conversions for config loading.
    """


class SmeltFileNotFoundError(SmeltError):
    """
    Raised when a file referenced by config is not found.
    """


class SmeltAlreadyNativeModule(SmeltError):
    """
    Raised when trying to convert to native an already native module.
    """


def path_exists(path: Path) -> TypeGuard[PathExists]:
    """
    Type magic.
    Allows representing existence of a given path in the type system.
    """
    return path.exists()


def assert_path_exists(path: str | Path) -> PathExists:
    path_obj = Path(path)
    if not path_exists(path_obj):
        raise SmeltConfigError(f"Following file/folder does not exist: {path}")
    return path_obj


def is_valid_import_path(import_path: str) -> TypeGuard[ImportPath]:
    """
    Checks if `import_path` is a valid import path in Python.
    """
    if not isinstance(import_path, str) or not import_path:
        return False

    parts = import_path.split(".")

    for part in parts:
        if not part.isidentifier():
            return False
        if keyword.iskeyword(part):
            return False
    return True


def assert_is_valid_import_path(import_path: str) -> ImportPath:
    if not is_valid_import_path(import_path):
        raise SmeltConfigError(f"Invalid python import path: {import_path}")
    return import_path


def is_valid_module_name(name: str) -> TypeGuard[ModuleName]:
    """
    Checks that `name` is a valid module name in Python.
    """
    if not is_valid_import_path(name):
        return False
    return "." not in name


def get_module_name(import_path: ImportPath) -> ModuleName:
    """
    Extracts the module name at the stem from `import_path`.
    """
    modname = import_path.split(".")[-1]
    assert is_valid_module_name(modname)
    return modname


class SmeltMissingModule(SmeltError): ...


def get_extension_suffix(target_triple: str) -> str:
    """
    Generate the C extension module filename.

    Parameters
    ----------
    target_triple: str
        The target triple, e.g., 'aarch64-linux-gnu'.

    Returns
    -------
    str
        The extension filename, e.g., '.cpython-312-aarch64-linux-gnu.so'
    """
    major = sys.version_info.major
    minor = sys.version_info.minor
    return f".cpython-{major}{minor}-{target_triple}.so"


def locate_module_by_import_path(import_path: ImportPath) -> PathExists:
    try:
        spec = importlib.util.find_spec(import_path)
    except ModuleNotFoundError as exc:
        # package not found
        raise SmeltMissingModule(
            f"{import_path} refers to some unknown packages in your environment. Is it installed ?"
        ) from exc

    if spec is None:  # module not found
        raise SmeltMissingModule(f"Module {import_path} not found. Is it installed ?")

    origin = spec.origin
    if origin is None:
        # PEP 420: path can be valid import without having an origin
        # if it refers to a namespace package,
        # as it might be found in multiple places.
        raise SmeltConfigError(
            f"Could not locate {import_path}. "
            "You might be referring to a namespace package, which is not a valid module to import"
        )
    module_path = Path(origin)
    assert path_exists(module_path), (
        "importlib.util.find_spec returned an invalid system path"
    )
    return module_path


@contextmanager
def import_shadowed_module(path: str) -> Generator[ModuleType, None, None]:
    """
    Allows importting a .py module shadowed by its .so counterpart
    (typically generated by mypy).
    .so python modules (generally, compiled C extensions) have higher priority on import
    over .py modules. This is is done so that module can provide optional C extensions
    with a pure Python implementation as fallback (e.g., if the interpreter is not CPython,
    or if the C extension cannot be compiled for the target platform)

    This function temporarily REMOVES the .so and proceeds to import the shadowed.
    It is a context manager and will restore the removed .so on exit,
    unless a new version has been built in the meantime (e.g., if the import
    was required as part of a recompile process)
    """
    with tempfile.TemporaryDirectory() as tmp:
        try:
            mod = importlib.import_module(path)
        except ImportError as exc:
            msg = f"Failed to import {path} while trying to mypycify"
            raise SmeltMissingModule(msg) from exc
        mod_path = mod.__file__
        assert mod_path is not None, f"Cannot module to mypycify: {path}"
        if mod_path.endswith((".py", ".pyc")):
            yield mod
        else:
            backup_path = os.path.join(tmp, os.path.basename(mod_path))
            shutil.copy(mod_path, backup_path)
            _logger.debug("Removing %s...", mod_path)
            os.remove(mod_path)
            os.sync()

            try:
                # using reload as otherwise importlib will use the cache
                # and re-import the now deleted .so
                mod = importlib.reload(mod)
                assert mod.__file__ is not None and not (mod.__file__.endswith(".so"))
                yield mod
            except ImportError as exc:
                msg = f"Failed to import {path} while trying to mypycify"
                raise SmeltMissingModule(msg) from exc


ImportPath = NewType("ImportPath", str)
FsPath = NewType("FsPath", str)


class ModpathType(Enum):
    """
    How locations of modules are passed:

    `IMPORT` means the module path is given as it would be used
    in a import statement.
    `FS` means the module is passed as a relative path.
    """

    IMPORT = auto()
    FS = auto()


@dataclass
class PathResolutionTrace:
    """
    Stores information about converting an import path to a filesystem path
    or vice-versa.
    """

    module_path: str
    import_path: str
    resolution_type: ModpathType

    def render(self) -> str:
        match self.resolution_type:
            case ModpathType.FS:
                return (
                    f"Resolving filesystem location of `{self.module_path}`:"
                    f"found `{self.import_path}`"
                )
            case ModpathType.IMPORT:
                return (
                    f"Resolving import path of file `{self.import_path}`:"
                    f"found `{self.module_path}`"
                )


def get_modpath_type(path: str) -> ModpathType:
    """
    Detects whether the passed type of the given `path`.
    """
    # TODO: stricter checks
    if path.endswith(".py") or path.endswith(".pyx"):
        return ModpathType.FS
    return ModpathType.IMPORT


class PackageLayout(Enum):
    """
    Which standard folder layout a given package is using.
    """

    SRC = auto()
    EPONYM = auto()


def detect_package_layout(package_name: str, package_root: str = ".") -> PackageLayout:
    """
    Tries detecting the package layout based on the file tree.
    """
    root = Path(package_root)
    has_src = (root / "src").exists()
    has_eponym_folder = (root / package_name).exists()

    if not has_src and not has_eponym_folder:
        raise RuntimeError(
            f"No `src` or `{package_name}` found cannot detect package layout"
        )
    if has_eponym_folder:
        if has_src:
            warnings.warn(
                "Package has both a `src` and a {package_name} folder at the root. "
                f"Assuming {package_name} contains source code, but that migth be wrong"
            )
            return PackageLayout.EPONYM

    if not (root / "src" / package_name).exists():
        raise RuntimeError(
            f"No `{package_name}` found inside `src` folder. "
            "If `src` is used as a substitute for your package name",
            "you need to configure that explicitly.",
        )
    return PackageLayout.SRC


def convert_to_path(import_path: ImportPath, file_extension: str = ".py") -> Path:
    *packages, module = import_path.split(".")
    module += file_extension
    return Path(*packages, module)


@overload
def find_mod_from_import_path_locally(
    import_path: ImportPath,
    file_extension: str,
    project_root: Path | None,
    *,
    should_exist: Literal[True],
) -> PathExists: ...
@overload
def find_mod_from_import_path_locally(
    import_path: ImportPath,
    file_extension: str,
    project_root: Path | None,
    *,
    should_exist: Literal[False],
) -> Path: ...


def find_mod_from_import_path_locally(
    import_path: ImportPath,
    file_extension: str = ".py",
    project_root: Path | None = None,
    *,
    should_exist: bool = True,
) -> PathExists | Path:
    """
    Given a module which import path is given by `import_path`,
    returns the module location as expected when iterating from the local working dir.
    """
    project_root = project_root or Path(".")
    path = project_root / convert_to_path(import_path, file_extension=file_extension)
    if not should_exist:
        return path
    if not path_exists(path):
        raise SmeltConfigError(
            f"Module {import_path} is not visible from current working dir:"
            f" {path} does not exists"
        )

    return path


@overload
def toggle_mod_path(path: str, to_type: Literal[ModpathType.IMPORT]) -> ImportPath: ...


@overload
def toggle_mod_path(path: str, to_type: Literal[ModpathType.FS]) -> FsPath: ...


def toggle_mod_path(
    path: str,
    to_type: ModpathType,
    src_folder: Path | None = None,
    exclude_root: bool = False,
) -> ImportPath | FsPath:
    """
    Changes the path to a module between 'import' style paths
    (a.b.c) and filesystem type (a/b/b).
    """
    src_folder = src_folder or Path(".")
    if get_modpath_type(path) == to_type:
        return cast(ImportPath | FsPath, path)

    match to_type:
        case ModpathType.IMPORT:
            path_obj = Path(path)
            resolved_path_obj = path_obj.resolve()
            resolved_src = src_folder.resolve()
            if not resolved_path_obj.is_relative_to(resolved_src):
                raise RuntimeError(f"`{src_folder}` is not a subpath of `path`")
            if exclude_root:
                resolved_path_obj = resolved_path_obj.parent
            rel_path = resolved_path_obj.relative_to(src_folder)
            modfile = rel_path.stem
            rel_path = rel_path.parent
            return cast(ImportPath, ".".join([*rel_path.parts, modfile]))

        case ModpathType.FS:
            path_chunks: list[str] = []
            package_or_namespace, *modules = path.split(".")
            if not modules:
                raise ValueError(
                    f"Modules import path only specifies the stem: {path}."
                    "You need to specify the whole import path"
                )
            if not exclude_root:
                path_chunks.append(package_or_namespace)
            *submodules, module = modules
            path_chunks.extend(submodules)
            module += ".py"
            path_chunks.append(module)
            final_path = src_folder.joinpath(*path_chunks)
            return cast(FsPath, str(final_path))

        case _ as unreachable:
            assert_never(unreachable)


def locate_module(
    mod_path: str,
    strategy: ModpathType = ModpathType.IMPORT,
    package_root: Path = Path("."),
) -> str:
    """
    Returns the full path to module at `mod_path`.
    If `mod_path` is passed a relative filepath, returns the absolute path.
    If it is passed as an import path ('.' seperated), tries importing and
    locates from the import"""
    match strategy:
        case ModpathType.IMPORT:
            import_path = locate_module_by_import_path(
                toggle_mod_path(mod_path, strategy)
            )
            if ctx := get_context():
                ctx.add_trace(
                    PathResolutionTrace(mod_path, mod_path, ModpathType.IMPORT)
                )
            return import_path

        case ModpathType.FS:
            fs_path = find_module_in_layout(
                toggle_mod_path(mod_path, strategy), package_root=package_root
            )
            if ctx := get_context():
                ctx.add_trace(PathResolutionTrace(fs_path, mod_path, ModpathType.FS))
            return fs_path

        case _ as unreachable:
            assert_never(unreachable)


def find_module_in_layout(mod_path: FsPath, package_root: str | None = None) -> str:
    """
    Given a module path on the filesystem,
    tries resolving that location within the package layout.
    """
    root = Path(package_root) if package_root else Path.cwd()
    src_location = root / "src"
    if src_location.exists():
        mod_potential_location = src_location / mod_path
        if mod_potential_location.exists():
            return str(mod_potential_location)

    mod_potential_location = root / mod_path
    if mod_potential_location.exists():
        return str(mod_potential_location)

    raise FileNotFoundError(f"Failed to locate {mod_path}")


@dataclass
class GenericExtension:
    """
    A wrapper on top of setuptools Extensions used by smelt.
    Used to provide metadata on how to manage these get_extension_suffixns in the install
    process, which is agnostic from the utility producing the extension in the first place.

    `import_path` keeps track of where the module should be placed in the final artifact.
    `src_path` keeps track of the original source files.
    `runtime` is optional and can be used to contain runtime utility module that's can
    be shared across modules (mypyc / Probably Nuitka in the future).
    """

    import_path: ImportPath
    src_path: PathExists
    dest_folder: Path
    extension: Extension
    runtime: Extension | None = None

    @classmethod
    def factory(
        cls,
        src_path: PathExists,
        import_path: ImportPath,
        extension: Extension | None = None,
        dest_folder: Path | None = None,
        runtime: Extension | None = None,
    ) -> Self:
        """
        Provides sound defaults for omitted fields.
        """
        extension = extension or Extension(
            name=get_module_name(import_path),
            sources=[src_path],
        )
        return cls(
            import_path=import_path,
            src_path=src_path,
            extension=extension,
            dest_folder=dest_folder or src_path.parent,
            runtime=runtime,
        )

    def get_dest_path(self, target_triple: str | None = None) -> Path:
        """
        Returns
        -------
        Path
            Full path for the final compiled .so file.
        """
        suffix = (
            sysconfig.get_config_var("EXT_SUFFIX")
            if target_triple is None
            else get_extension_suffix(target_triple)
        )
        ext_so_name = f"{self.name}{suffix}"
        return self.dest_folder / ext_so_name

    @property
    def name(self) -> ModuleName:
        """
        Returns the module name that given to the extension.
        """
        return get_module_name(self.import_path)

    def get_runtime_dest_path(self, target_triple: str | None = None) -> Path:
        """
        Returns
        -------
        Path
            Full path for the final runtime .so file.
        """
        suffix = (
            sysconfig.get_config_var("EXT_SUFFIX")
            if target_triple is None
            else get_extension_suffix(target_triple)
        )
        ext_so_name = f"{self.name}__mypyc{suffix}"
        return self.dest_folder / ext_so_name


class PackageRootPath(NamedTuple):
    module_import_path: ImportPath
    path: PathExists

    @classmethod
    def from_path(cls, path: PathExists) -> Self:
        package_name = path.parts[-1]
        if not is_valid_import_path(package_name):
            raise SmeltConfigError(
                f"Module path uses invalid characters: {package_name}"
            )
        return cls(package_name, path)


@dataclass
class PathSolver:
    known_roots: list[PackageRootPath] = field(default_factory=list)
    project_root: Path = Path(".")

    @classmethod
    def from_installed_import_paths(cls, *import_path: ImportPath):
        """
        Creates a PathSolver for the given import_paths, given that the reference modules
        are installed in the local environement.
        This relies on side-effects (modules location relies on specs given by importlib)
            and not suitable for build-time usage,
            but can be used as a shortcut for usage within an already installed environement
            (i.e. frontend usage).
        """
        roots = [
            PackageRootPath(p, locate_module_by_import_path(p)) for p in import_path
        ]
        return cls(roots)

    @overload
    def resolve_import_path(
        self,
        import_path: ImportPath,
        file_extension: str,
        *,
        should_exist: Literal[True],
    ) -> PathExists: ...
    @overload
    def resolve_import_path(
        self,
        import_path: ImportPath,
        file_extension: str,
        *,
        should_exist: Literal[False],
    ) -> Path: ...

    def resolve_import_path(
        self,
        import_path: ImportPath,
        file_extension: str = ".py",
        *,
        should_exist: bool = True,
    ) -> Path | PathExists:
        # TODO: detect when target module is already native
        for root_path, path in self.known_roots:
            if import_path == root_path:
                return path
            if not import_path.startswith(root_path):
                continue
            root = Path(path)
            sub_import_path = import_path.replace(root_path + ".", "")
            assert is_valid_import_path(sub_import_path), (
                f"internal path substitution logic error {sub_import_path}"
            )

            subpath = convert_to_path(sub_import_path)
            module_path = root / subpath
            if should_exist and not path_exists(module_path):
                raise SmeltConfigError(
                    f"Module {import_path} should resolve to this location "
                    f"but we did not find it: {module_path}"
                )
            return module_path

        return find_mod_from_import_path_locally(
            import_path,
            file_extension=file_extension,
            project_root=self.project_root,
            should_exist=should_exist,
        )
