import importlib.util
import logging
import sys
from enum import StrEnum
from pathlib import Path
from types import ModuleType

from faebryk.libs.util import (
    ConfigFlag,
    ConfigFlagEnum,
    FileChangedWatcher,
    debug_perf,
    get_python_lib,
    global_lock,
    is_editable_install,
    run_live,
)

logger = logging.getLogger(__name__)

_thisfile = Path(__file__)
_thisdir = _thisfile.parent
_build_dir = _thisdir / "zig-out" / "lib"
_pyi_dir = _build_dir / "gen"


class ReleaseMode(StrEnum):
    DEBUG = "Debug"
    """Optimizations off and safety on (default)"""
    RELEASE_SAFE = "ReleaseSafe"
    """Optimizations on and safety on"""
    RELEASE_FAST = "ReleaseFast"
    """Optimizations on and safety off"""
    RELEASE_SMALL = "ReleaseSmall"
    """Size optimizations on and safety off"""


RELEASEMODE = ConfigFlagEnum(
    "ZIG_RELEASEMODE",
    default=ReleaseMode.RELEASE_FAST,
    enum=ReleaseMode,
    descr="https://ziglang.org/documentation/0.15.2/#Build-Mode",
)
NORECOMPILE = ConfigFlag("ZIG_NORECOMPILE", default=False)

_source_hash_file = _build_dir / ".zig-source-hash"


def _extension_suffix() -> str:
    return ".pyd" if sys.platform == "win32" else ".so"


def _compute_source_hash() -> str:
    """Hash all zig source files + build config to detect changes."""
    from .compute_hash import compute_source_hash

    return compute_source_hash(zig_dir=_thisdir, release_mode=RELEASEMODE.value)


def _zig_sources_changed() -> bool:
    """Check if zig sources changed since last successful build."""
    ext = _extension_suffix()
    if not (_build_dir / f"pyzig{ext}").exists():
        return True
    if not (_build_dir / f"pyzig_sexp{ext}").exists():
        return True
    if not _source_hash_file.exists():
        return True
    return _source_hash_file.read_text().strip() != _compute_source_hash()


def _save_source_hash() -> None:
    _source_hash_file.write_text(_compute_source_hash())


@debug_perf
def compile_zig():
    from faebryk.libs.header import get_header
    from faebryk.libs.util import insert_into_file, run_processes

    zig_cmd = [sys.executable, "-m", "ziglang"]
    python_lib = get_python_lib()

    pyi_watcher = FileChangedWatcher(
        _pyi_dir, method=FileChangedWatcher.CheckMethod.HASH, glob="*/*.pyi"
    )

    # Use zig build system with Python configuration
    cmd = [
        *zig_cmd,
        "build",
        "python-ext",
        f"-Doptimize={RELEASEMODE.value}",
        f"-Dpython-include={python_lib.include_path}",
        f"-Dpython-lib={python_lib.name}",
        *([f"-Dpython-lib-dir={python_lib.dir_path}"] if python_lib.dir_path else []),
    ]

    logger.info(f"Building with command: {' '.join(cmd)}")
    with global_lock(_build_dir / "lock", timeout_s=60):
        run_live(cmd, cwd=Path(__file__).parent)

    # Process all generated pyi files
    # only if any change in the pyi, formatting is expensive
    if changed := pyi_watcher.has_changed():
        logger.info(f"Changed pyi files: {changed}")
        header = (
            get_header()
            + "\n"
            + "# This file is auto-generated by zig.\n"
            + "# Do not edit this file directly; edit the corresponding\n"
            + "# Zig file instead.\n\n"
        )

        def _out_for_in(path: Path) -> Path:
            return _thisdir / path.relative_to(_build_dir)

        for deleted in changed.deleted:
            _out_for_in(deleted).unlink(missing_ok=True)

        pyi_outs = [
            (pyi_source, _out_for_in(pyi_source))
            for pyi_source in changed.changed + changed.created
        ]
        for pyi_source, pyi_out in pyi_outs:
            pyi_out.parent.mkdir(parents=True, exist_ok=True)
            insert_into_file(pyi_source, header, pyi_out)

        run_processes(
            [
                [sys.executable, "-m", "black", "--pyi", "--fast", str(pyi_out)]
                for _, pyi_out in pyi_outs
            ]
        )
        run_processes(
            [
                [sys.executable, "-m", "ruff", "check", "--fix", str(pyi_out)]
                for _, pyi_out in pyi_outs
            ]
        )


def _load_module(path: list[str], module: ModuleType):
    _module_path = ".".join(path)
    sys.modules[_module_path] = module
    for name in dir(module):
        if name.startswith("_"):
            continue
        obj = getattr(module, name)
        path_ = path + [name]
        if isinstance(obj, ModuleType):
            _load_module(path_, obj)


def _load_ext(dir: Path, name: str) -> ModuleType:
    path = dir / f"{name}{_extension_suffix()}"

    spec = importlib.util.spec_from_file_location(name, path)
    assert spec and spec.loader
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


def load():
    return _load_ext(_build_dir, "pyzig")


def load_sexp():
    return _load_ext(_build_dir, "pyzig_sexp")


# Fallback to editable build-on-import
if is_editable_install():
    if not NORECOMPILE.get() and _zig_sources_changed():
        compile_zig()
        _save_source_hash()
    pyzig = load()

    # load sexp
    sexp = load_sexp()
    # Attach the sexp module under pyzig.gen to keep the old API surface.
    gen = getattr(pyzig, "gen", None)
    if gen is not None:
        setattr(gen, "sexp", sexp)
    # Populate sys.modules entries for nested import paths.
    sys.modules["pyzig.gen.sexp"] = sexp
    sys.modules["faebryk.core.zig.gen.sexp"] = sexp
    sexp.__name__ = "pyzig.gen.sexp"

    _load_module(["faebryk", "core", "zig"], pyzig)
else:
    import pyzig
    import pyzig_sexp

    # Attach the sexp module under pyzig.gen to keep the old API surface.
    gen = getattr(pyzig, "gen", None)
    if gen is not None:
        setattr(gen, "sexp", pyzig_sexp)
    sys.modules["pyzig.gen.sexp"] = pyzig_sexp
    sys.modules["faebryk.core.zig.gen.sexp"] = pyzig_sexp
    pyzig_sexp.__name__ = "pyzig.gen.sexp"

    _load_module(["faebryk", "core", "zig"], pyzig)
