from __future__ import annotations

import importlib
import inspect
import os
import pathlib
import sys
import tempfile
from typing import List

import pkg_resources
from klvm_tools_rs import compile_klvm as compile_klvm_rust

from chik.types.blockchain_format.program import Program
from chik.types.blockchain_format.serialized_program import SerializedProgram
from chik.util.lock import Lockfile

compile_klvm_py = None

recompile_requested = (
    (os.environ.get("CHIK_DEV_COMPILE_KLVM_ON_IMPORT", "") != "") or ("pytest" in sys.modules)
) and os.environ.get("CHIK_DEV_COMPILE_KLVM_DISABLED", None) is None


def translate_path(p_):
    p = str(p_)
    if os.path.isdir(p):
        return p
    else:
        module_object = importlib.import_module(p)
        return os.path.dirname(inspect.getfile(module_object))


# Handle optional use of python klvm_tools if available and requested
if "KLVM_TOOLS" in os.environ:
    try:
        from klvm_tools.klvmc import compile_klvm as compile_klvm_py_candidate

        compile_klvm_py = compile_klvm_py_candidate
    finally:
        pass


def compile_klvm_in_lock(full_path: pathlib.Path, output: pathlib.Path, search_paths: List[pathlib.Path]):
    # Compile using rust (default)

    # Ensure path translation is done in the idiomatic way currently
    # expected.  It can use either a filesystem path or name a python
    # module.
    treated_include_paths = list(map(translate_path, search_paths))
    res = compile_klvm_rust(str(full_path), str(output), treated_include_paths)

    if "KLVM_TOOLS" in os.environ and os.environ["KLVM_TOOLS"] == "check" and compile_klvm_py is not None:
        # Simple helper to read the compiled output
        def sha256file(f):
            import hashlib

            m = hashlib.sha256()
            m.update(open(f).read().strip().encode("utf8"))
            return m.hexdigest()

        orig = "%s.orig" % output

        compile_klvm_py(full_path, orig, search_paths=search_paths)
        orig256 = sha256file(orig)
        rs256 = sha256file(output)

        if orig256 != rs256:
            print("Compiled original %s: %s vs rust %s\n" % (full_path, orig256, rs256))
            print("Aborting compilation due to mismatch with rust")
            assert orig256 == rs256
        else:
            print("Compilation match %s: %s\n" % (full_path, orig256))

    return res


def compile_klvm(full_path: pathlib.Path, output: pathlib.Path, search_paths: List[pathlib.Path] = []):
    with Lockfile.create(pathlib.Path(tempfile.gettempdir()) / "klvm_compile" / full_path.name):
        compile_klvm_in_lock(full_path, output, search_paths)


def load_serialized_klvm(
    klvm_filename, package_or_requirement=__name__, include_standard_libraries: bool = False, recompile: bool = True
) -> SerializedProgram:
    """
    This function takes a .clsp file in the given package and compiles it to a
    .clsp.hex file if the .hex file is missing or older than the .clsp file, then
    returns the contents of the .hex file as a `Program`.

    klvm_filename: file name
    package_or_requirement: usually `__name__` if the klvm file is in the same package
    """
    hex_filename = f"{klvm_filename}.hex"

    # Set the CHIK_DEV_COMPILE_KLVM_ON_IMPORT environment variable to anything except
    # "" or "0" to trigger automatic recompilation of the Chiklisp on load.
    if recompile:
        try:
            if pkg_resources.resource_exists(package_or_requirement, klvm_filename):
                # Establish whether the size is zero on entry
                full_path = pathlib.Path(pkg_resources.resource_filename(package_or_requirement, klvm_filename))
                output = full_path.parent / hex_filename
                if not output.exists() or os.stat(full_path).st_mtime > os.stat(output).st_mtime:
                    search_paths = [full_path.parent]
                    if include_standard_libraries:
                        # we can't get the dir, but we can get a file then get its parent.
                        chik_puzzles_path = pathlib.Path(
                            pkg_resources.resource_filename(__name__, "__init__.py")
                        ).parent
                        search_paths.append(chik_puzzles_path)
                    compile_klvm(full_path, output, search_paths=search_paths)

        except NotImplementedError:
            # pyinstaller doesn't support `pkg_resources.resource_exists`
            # so we just fall through to loading the hex klvm
            pass

    klvm_hex = pkg_resources.resource_string(package_or_requirement, hex_filename).decode("utf8")
    assert len(klvm_hex.strip()) != 0
    klvm_blob = bytes.fromhex(klvm_hex)
    return SerializedProgram.from_bytes(klvm_blob)


def load_klvm(
    klvm_filename,
    package_or_requirement=__name__,
    include_standard_libraries: bool = False,
    recompile: bool = True,
) -> Program:
    return Program.from_bytes(
        bytes(
            load_serialized_klvm(
                klvm_filename,
                package_or_requirement=package_or_requirement,
                include_standard_libraries=include_standard_libraries,
                recompile=recompile,
            )
        )
    )


def load_klvm_maybe_recompile(
    klvm_filename,
    package_or_requirement=__name__,
    include_standard_libraries: bool = False,
    recompile: bool = recompile_requested,
) -> Program:
    return load_klvm(
        klvm_filename=klvm_filename,
        package_or_requirement=package_or_requirement,
        include_standard_libraries=include_standard_libraries,
        recompile=recompile,
    )


def load_serialized_klvm_maybe_recompile(
    klvm_filename,
    package_or_requirement=__name__,
    include_standard_libraries: bool = False,
    recompile: bool = recompile_requested,
) -> SerializedProgram:
    return load_serialized_klvm(
        klvm_filename=klvm_filename,
        package_or_requirement=package_or_requirement,
        include_standard_libraries=include_standard_libraries,
        recompile=recompile,
    )