import modulegraph.modulegraph  # type: ignore
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys
import time
import site
from typing import Any, Callable

import basilisk
from basilisk.build_wasm_binary_or_exit import build_wasm_binary_or_exit
from basilisk.cargotoml import generate_cargo_toml, generate_cargo_lock
from basilisk.colors import red, green, dim
from basilisk.run_basilisk_generate_or_exit import run_basilisk_generate_or_exit
from basilisk.timed import timed, timed_inline
from basilisk.types import Args, Paths


@timed
def main():
    # TODO this way of installing the extension is just temporary
    # TODO we should use the official dfx extension install command
    # TODO are the dfx extensions repository once those mature
    if sys.argv[1] == "install-dfx-extension":
        subprocess.run(
            ["./install.sh"],
            cwd=os.path.join(
                os.path.dirname(basilisk.__file__), "compiler", "dfx_extension"
            ),
        )
        return

    args = parse_args_or_exit(sys.argv)
    paths = create_paths(args)
    is_verbose = args["flags"]["verbose"] or os.environ.get("BASILISK_VERBOSE") == "true"

    subprocess.run(
        [
            f"{paths['compiler']}/install_rust_dependencies.sh",
            basilisk.__version__,
            basilisk.__rust_version__,
        ]
    )

    # This is the name of the canister passed into python -m basilisk from the dfx.json build command
    canister_name = args["canister_name"]

    verbose_mode_qualifier = " in verbose mode" if is_verbose else ""

    print(f"\nBuilding canister {green(canister_name)}{verbose_mode_qualifier}\n")

    python_backend = os.environ.get("BASILISK_PYTHON_BACKEND", "cpython")

    # CPython template mode (default): skip Rust codegen entirely, just bundle Python + manipulate wasm
    if python_backend == "cpython":
        # Only need the canister staging dir (for output wasm) and Python source
        os.makedirs(paths["canister"], exist_ok=True)

        cargo_env = {
            **os.environ.copy(),
            "CARGO_TARGET_DIR": paths["global_basilisk_target_dir"],
            "CARGO_HOME": paths["global_basilisk_rust_dir"],
            "RUSTUP_HOME": paths["global_basilisk_rust_dir"],
        }

        # Bundle the user's Python code (no Rust codegen needed)
        bundle_python_code(paths)

        build_wasm_binary_or_exit(
            paths,
            canister_name,
            cargo_env,
            verbose=is_verbose,
            label=f"[1/1] ⚡ Building Wasm from template...",
        )

        print(f"\n🎉 Built canister {green(canister_name)} at {dim(paths['wasm'])}")
        return

    # Standard mode: generate Rust code and compile
    # Copy all of the Rust project structure from the pip package to an area designed for Rust compiling
    if os.path.exists(paths["canister"]):
        shutil.rmtree(paths["canister"])
    shutil.copytree(paths["compiler"], paths["canister"], dirs_exist_ok=True)
    create_file(f"{paths['canister']}/Cargo.toml", generate_cargo_toml(canister_name, python_backend))
    create_file(f"{paths['canister']}/Cargo.lock", generate_cargo_lock())

    # Add CARGO_TARGET_DIR to env for all cargo commands
    cargo_env = {
        **os.environ.copy(),
        "CARGO_TARGET_DIR": paths["global_basilisk_target_dir"],
        "CARGO_HOME": paths["global_basilisk_rust_dir"],
        "RUSTUP_HOME": paths["global_basilisk_rust_dir"],
    }
    if not os.path.exists(paths["global_basilisk_bin_dir"]):
        os.makedirs(paths["global_basilisk_bin_dir"])

    compile_python_or_exit(
        paths, cargo_env, verbose=is_verbose, label="[1/2] 🔨 Compiling Python..."
    )

    # Fix panic!(err) patterns in generated code (cdk_framework uses deprecated syntax)
    fixup_generated_code(paths)

    build_wasm_binary_or_exit(
        paths,
        canister_name,
        cargo_env,
        verbose=is_verbose,
        label=f"[2/2] 🚧 Building Wasm binary...",
    )

    print(f"\n🎉 Built canister {green(canister_name)} at {dim(paths['wasm'])}")


def parse_args_or_exit(args: list[str]) -> Args:
    args = args[1:]  # Discard the path to basilisk

    flags = [arg for arg in args if (arg.startswith("-") or arg.startswith("--"))]
    args = [arg for arg in args if not (arg.startswith("-") or arg.startswith("--"))]

    if len(args) == 0:
        print(f"\nbasilisk {basilisk.__version__}")
        print("\nUsage: basilisk [-v|--verbose] <canister_name> <entry_point>")
        sys.exit(0)

    if len(args) != 2:
        print(red("\n💣 Basilisk error: wrong number of arguments\n"))
        print("Usage: basilisk [-v|--verbose] <canister_name> <entry_point>")
        print("\n💀 Build failed!")
        sys.exit(1)

    return {
        "empty": False,
        "flags": {"verbose": "--verbose" in flags or "-v" in flags},
        "canister_name": args[0],
        "entry_point": args[1],
    }


def create_paths(args: Args) -> Paths:
    canister_name = args["canister_name"]

    # This is the path to the developer's entry point Python file passed into python -m basilisk from the dfx.json build command
    py_entry_file_path = args["entry_point"]

    # This is the Python module name of the developer's Python project, derived from the entry point Python file passed into python -m basilisk from the dfx.json build command
    py_entry_module_name = Path(py_entry_file_path).stem

    # This is the location of all code used to generate the final canister Rust code
    canister_path = f".basilisk/{canister_name}"

    # We want to bundle/gather all Python files into the python_source directory for RustPython freezing
    # The location that Basilisk will look to when running py_freeze!
    # py_freeze! will compile all of the Python code in the directory recursively (modules must have an __init__.py to be included)
    python_source_path = f"{canister_path}/python_source"

    py_file_names_file_path = f"{canister_path}/py_file_names.csv"

    # This is the path to the developer's Candid file as resolved by dfx
    did_path = os.environ.get("CANISTER_CANDID_PATH")

    if did_path is None:
        raise Exception("Basilisk: CANISTER_CANDID_PATH is not defined")

    # This is the path to the Basilisk compiler Rust code delivered with the Python package
    compiler_path = os.path.dirname(basilisk.__file__) + "/compiler"

    # This is the final generated Rust file that is the canister
    lib_path = f"{canister_path}/src/lib.rs"

    # This is the location of the Candid file generated from the final generated Rust file
    generated_did_path = f"{canister_path}/index.did"

    # This is the unzipped generated Wasm that is the canister
    wasm_path = f"{canister_path}/{canister_name}.wasm"

    # This is where we store custom Python modules, such as stripped-down versions of stdlib modules
    custom_modules_path = f"{compiler_path}/custom_modules"

    home_dir = os.path.expanduser("~")
    global_basilisk_config_dir = f"{home_dir}/.config/basilisk"
    global_basilisk_version_dir = f"{global_basilisk_config_dir}/{basilisk.__version__}"
    global_basilisk_rust_dir = f"{global_basilisk_config_dir}/rust/{basilisk.__rust_version__}"
    global_basilisk_rust_bin_dir = f"{global_basilisk_rust_dir}/bin"
    global_basilisk_target_dir = f"{global_basilisk_config_dir}/rust/target"
    global_basilisk_bin_dir = f"{global_basilisk_config_dir}/{basilisk.__version__}/bin"

    return {
        "py_entry_file": py_entry_file_path,
        "py_entry_module_name": py_entry_module_name,
        "canister": canister_path,
        "python_source": python_source_path,
        "py_file_names_file": py_file_names_file_path,
        "did": did_path,
        "compiler": compiler_path,
        "lib": lib_path,
        "generated_did": generated_did_path,
        "wasm": wasm_path,
        "custom_modules": custom_modules_path,
        "global_basilisk_config_dir": global_basilisk_config_dir,
        "global_basilisk_version_dir": global_basilisk_version_dir,
        "global_basilisk_rust_dir": global_basilisk_rust_dir,
        "global_basilisk_rust_bin_dir": global_basilisk_rust_bin_dir,
        "global_basilisk_target_dir": global_basilisk_target_dir,
        "global_basilisk_bin_dir": global_basilisk_bin_dir,
    }


@timed_inline
def compile_python_or_exit(
    paths: Paths, cargo_env: dict[str, str], verbose: bool = False
):
    bundle_python_code(paths)
    run_basilisk_generate_or_exit(paths, cargo_env, verbose)
    run_rustfmt_or_exit(paths, cargo_env, verbose)


def bundle_python_code(paths: Paths):
    # Begin module bundling/gathering process
    path = (
        list(filter(lambda x: x.startswith(os.getcwd()), sys.path))
        + [p for p in sys.path if "site-packages" in p or "dist-packages" in p]
        + [
            os.path.dirname(paths["py_entry_file"]),
            os.path.dirname(os.path.dirname(basilisk.__file__)),
        ]
        + site.getsitepackages()
        + [site.getusersitepackages()]
    )

    graph = modulegraph.modulegraph.ModuleGraph(path)  # type: ignore
    entry_point = graph.run_script(paths["py_entry_file"])  # type: ignore

    python_source_path = paths["python_source"]

    if os.path.exists(python_source_path):
        shutil.rmtree(python_source_path)

    os.makedirs(python_source_path)

    # Copy our custom Python modules into the python_source directory
    shutil.copytree(paths["custom_modules"], python_source_path, dirs_exist_ok=True)

    flattened_graph = list(graph.flatten(start=entry_point))  # type: ignore

    for node in flattened_graph:  # type: ignore
        if type(node) == modulegraph.modulegraph.Script:  # type: ignore
            shutil.copy(
                node.filename, f"{python_source_path}/{os.path.basename(node.filename)}"  # type: ignore
            )

        if type(node) == modulegraph.modulegraph.SourceModule:  # type: ignore
            # Use dotted identifier to build path so modules with the same
            # basename (e.g. _pytest.main vs the entry-point main) don't
            # overwrite each other.
            dest_name = node.identifier.replace(".", os.sep) + ".py"  # type: ignore
            dest_path = f"{python_source_path}/{dest_name}"
            os.makedirs(os.path.dirname(dest_path), exist_ok=True)
            shutil.copy(node.filename, dest_path)  # type: ignore

        if type(node) == modulegraph.modulegraph.Package:  # type: ignore
            # Skip the installed basilisk package (compiler, shell, etc.)
            # but allow canister-side subpackages like basilisk.os
            if should_skip_package(node.identifier, node.packagepath[0]):  # type: ignore
                continue
            # Convert dotted identifier to path (e.g. basilisk.os -> basilisk/os)
            dest_dir = node.identifier.replace(".", os.sep)  # type: ignore
            shutil.copytree(
                node.packagepath[0],  # type: ignore
                f"{python_source_path}/{dest_dir}",  # type: ignore
                dirs_exist_ok=True,
                ignore=ignore_specific_dir,
            )
            # Ensure parent __init__.py files exist for nested packages
            _ensure_parent_inits(python_source_path, dest_dir)

        if type(node) == modulegraph.modulegraph.NamespacePackage:  # type: ignore
            dest_dir = node.identifier.replace(".", os.sep)  # type: ignore
            shutil.copytree(
                node.packagepath[0],  # type: ignore
                f"{python_source_path}/{dest_dir}",  # type: ignore
                dirs_exist_ok=True,
                ignore=ignore_specific_dir,
            )
            _ensure_parent_inits(python_source_path, dest_dir)

    py_file_names = list(  # type: ignore
        filter(
            lambda filename: filename is not None and filename.endswith(".py"),  # type: ignore
            map(
                lambda node: node.filename,  # type: ignore
                filter(
                    lambda node: node.filename  # type: ignore
                    is not "-",  # This filters out namespace packages
                    flattened_graph,  # type: ignore
                ),  # type: ignore
            ),  # type: ignore
        )  # type: ignore
    )

    create_file(paths["py_file_names_file"], ",".join(py_file_names))  # type: ignore


def ignore_specific_dir(dirname: str, filenames: list[str]) -> list[str]:
    if "basilisk_post_install/src/Lib" in dirname:
        return filenames
    else:
        return []


def should_skip_package(node_identifier: str, node_packagepath: str) -> bool:
    """Skip the top-level installed basilisk package (compiler, shell, etc.)
    but allow canister-side subpackages like basilisk.os through."""
    if "site-packages" not in node_packagepath:
        return False
    # Skip the top-level basilisk package but not its subpackages
    if node_identifier == "basilisk":
        return True
    return False


def _ensure_parent_inits(python_source_path: str, dest_dir: str):
    """Create empty __init__.py in parent directories for nested packages.

    E.g. for dest_dir='basilisk/os', ensures basilisk/__init__.py exists
    so that 'from basilisk.os import ...' works at runtime.
    """
    parts = dest_dir.split(os.sep)
    for i in range(1, len(parts)):
        parent = os.path.join(python_source_path, *parts[:i])
        init_file = os.path.join(parent, "__init__.py")
        if not os.path.exists(init_file):
            os.makedirs(parent, exist_ok=True)
            with open(init_file, "w") as f:
                f.write("")


def parse_basilisk_generate_error(stdout: bytes) -> str:
    err = stdout.decode("utf-8")
    std_err_lines = err.splitlines()
    try:
        line_where_error_message_starts = next(
            i
            for i, v in enumerate(std_err_lines)
            if v.startswith("thread 'main' panicked at '")
        )
        line_where_error_message_ends = next(
            i for i, v in enumerate(std_err_lines) if "', src/" in v
        )
    except:
        return err

    err_lines = std_err_lines[
        line_where_error_message_starts : line_where_error_message_ends + 1
    ]
    err_lines[0] = err_lines[0].replace("thread 'main' panicked at '", "")
    err_lines[-1] = re.sub("', src/.*", "", err_lines[-1])

    return red("\n".join(err_lines))


def run_rustfmt_or_exit(paths: Paths, cargo_env: dict[str, str], verbose: bool = False):
    rustfmt_result = subprocess.run(
        [
            f"{paths['global_basilisk_rust_bin_dir']}/rustfmt",
            "--edition=2018",
            paths["lib"],
        ],
        capture_output=not verbose,
        env=cargo_env,
    )

    if rustfmt_result.returncode != 0:
        print(red("\n💣 Basilisk error: internal Rust formatting"))
        print(
            f'\nPlease open an issue at https://github.com/smart-social-contracts/basilisk/issues/new\nincluding this message and the following error:\n\n {red(rustfmt_result.stderr.decode("utf-8"))}'
        )
        print("💀 Build failed")
        sys.exit(1)


def fixup_generated_code(paths: Paths):
    """Fix known issues in generated Rust code before compilation."""
    lib_path = paths["lib"]
    if not os.path.exists(lib_path):
        return
    with open(lib_path, "r") as f:
        content = f.read()
    # cdk_framework generates panic!(err) which is invalid in newer Rust editions;
    # the err type is (RejectionCode, String) which only implements Debug, not Display
    content = re.sub(r'panic!\(err\)', 'panic!("{:?}", err)', content)
    with open(lib_path, "w") as f:
        f.write(content)


def create_file(file_path: str, contents: str):
    file = open(file_path, "w")
    file.write(contents)
    file.close()


def inline_timed(
    label: str,
    body: Callable[..., Any],
    *args: Any,
    verbose: bool = False,
    **kwargs: Any,
) -> float:
    print(label)
    start_time = time.time()
    body(*args, verbose=verbose, **kwargs)
    end_time = time.time()
    duration = end_time - start_time

    if verbose:
        print(f"{label} finished in {round(duration, 2)}s")
    else:
        move_cursor_up_one_line = "\x1b[1A"
        print(f'{move_cursor_up_one_line}{label} {dim(f"{round(duration, 2)}s")}')

    return end_time - start_time


main()
