from __future__ import annotations

import os
import shutil
from pathlib import Path
from typing import Optional, cast

import click
from loguru import logger
from tomlkit.api import document, dumps, inline_table, nl, table
from tomlkit.items import Array, KeyType, SingleKey, Trivia
from tqdm.auto import tqdm

from sbt.config import PBTConfig
from sbt.misc import (
    ExecProcessError,
    IncompatibleDependencyError,
    NewEnvVar,
    exec,
    mask_file,
    venv_path,
)
from sbt.package.discovery import (
    discover_packages,
    parse_pep518_pyproject,
    parse_version_spec,
)
from sbt.package.graph import PkgGraph
from sbt.package.package import DepConstraint, DepConstraints, Package

# environment variables that will be passed to the subprocess
PASSTHROUGH_ENVS = [
    "PATH",
    "CC",
    "CXX",
    "LIBCLANG_PATH",
    "LD_LIBRARY_PATH",
    "DYLD_LIBRARY_PATH",
    "C_INCLUDE_PATH",
    "CPLUS_INCLUDE_PATH",
    "HOME",
    "CARGO_HOME",
    "RUSTUP_HOME",
]


@click.command()
@click.argument("package", default="")
@click.option("--cwd", default=".", help="Override current working directory")
@click.option(
    "--ignore-invalid-pkg",
    is_flag=True,
    help="whether to ignore invalid packages",
)
@click.option(
    "--ignore-invalid-dependency",
    is_flag=True,
    help="whether to ignore invalid dependencies",
)
@click.option(
    "--all-packages",
    is_flag=True,
    help="Install all discovered packages as dependencies of the target one",
)
@click.option(
    "-v",
    "--verbose",
    is_flag=True,
    help="increase verbosity",
)
def install(
    package: str,
    cwd: str = ".",
    ignore_invalid_pkg: bool = False,
    ignore_invalid_dependency: bool = False,
    all_packages: bool = False,
    verbose: bool = False,
):
    """Install a package and its local dependencies in editable mode"""
    cwd = os.path.abspath(cwd)
    cfg = PBTConfig.from_dir(cwd)

    # discovery packages
    packages = discover_packages(
        cfg.cwd,
        cfg.cache_dir,
        cfg.ignore_directories,
        cfg.ignore_directory_names,
        ignore_invalid_package=ignore_invalid_pkg,
    )
    loc2pkg = {pkg.location: pkg for pkg in packages.values()}
    if package == "":
        # use the package located in the current directory
        if cfg.cwd in loc2pkg:
            package = loc2pkg[cfg.cwd].name
        else:
            raise ValueError(
                f"Cannot find a package in the current directory {cfg.cwd}"
            )

    local_dep_pkgs = None
    if all_packages:
        local_dep_pkgs = [pkg for pkg in packages.values() if pkg.name != package]
    install_pkg(
        packages[package], packages, cfg, ignore_invalid_dependency, local_dep_pkgs
    )


@click.command()
@click.argument("dependency")
@click.option("--package", default="", help="The target package to add dependency to")
@click.option("--cwd", default=".", help="Override current working directory")
@click.option(
    "--no-dep-dep", is_flag=True, help="Do not install dependencies of the dependency"
)
@click.option(
    "--ignore-invalid-pkg",
    is_flag=True,
    help="whether to ignore invalid packages",
)
@click.option(
    "--ignore-invalid-dependency",
    is_flag=True,
    help="whether to ignore invalid dependencies",
)
@click.option(
    "-v",
    "--verbose",
    is_flag=True,
    help="increase verbosity",
)
def add(
    dependency: str,
    package: str = "",
    cwd: str = ".",
    no_dep_dep: bool = False,
    ignore_invalid_pkg: bool = False,
    ignore_invalid_dependency: bool = False,
    verbose: bool = False,
):
    """Add a local package as a dependency to the target package"""
    cwd = os.path.abspath(cwd)
    cfg = PBTConfig.from_dir(cwd)

    # discovery packages
    packages = discover_packages(
        cfg.cwd,
        cfg.cache_dir,
        cfg.ignore_directories,
        cfg.ignore_directory_names,
        ignore_invalid_package=ignore_invalid_pkg,
    )
    loc2pkg = {pkg.location: pkg for pkg in packages.values()}
    if package == "":
        # use the package located in the current directory
        if cfg.cwd in loc2pkg:
            package = loc2pkg[cfg.cwd].name
        else:
            raise ValueError(
                f"Cannot find a package in the current directory {cfg.cwd}"
            )

    # step 1: get the dependency package and add its to the list of discover packages
    if dependency not in packages:
        # dependency is a path
        deppkg = parse_pep518_pyproject(Path(dependency))
        packages[deppkg.name] = deppkg
    else:
        deppkg = packages[dependency]

    # step 2: update the list of manually installed
    pkg_cache_dir = cfg.pkg_cache_dir(packages[package])
    manual_deps = packages[package].find_manually_installed_dependencies(pkg_cache_dir)
    manual_deps.append(deppkg.location)
    packages[package].save_manually_installed_dependencies(
        pkg_cache_dir, sorted(set(manual_deps))
    )

    # step 3: gather dependencies of the package -- include the one from manually installed packages
    if no_dep_dep:
        pkg_venv_path = venv_path(
            packages[package].name,
            packages[package].location,
            cfg.python_virtualenv_path,
            cfg.get_python_path(),
        )
        install_bare_pkg(packages[package], cfg, pkg_venv_path)
    else:
        install_pkg(
            packages[package], packages, cfg, ignore_invalid_dependency, [deppkg]
        )


def install_pkg(
    target_pkg: Package,
    packages: dict[str, Package],
    cfg: PBTConfig,
    ignore_invalid_dependency: bool,
    local_dep_pkgs: Optional[list[Package]],
):
    loc2pkg = {pkg.location: pkg for pkg in packages.values()}

    # step 0: gather dependencies of the package -- include the one from manually installed packages
    target_pkg.dependencies.update(
        {
            (p := loc2pkg[loc]).name: [
                DepConstraint(version_spec="=={}".format(p.version))
            ]
            for loc in target_pkg.find_manually_installed_dependencies(
                cfg.pkg_cache_dir(target_pkg)
            )
        }
    )
    pkg_graph = PkgGraph.from_pkgs(packages)

    thirdparty_pkgs: dict[str, tuple[set[str], DepConstraints]] = {}
    invalid_thirdparty_pkgs: set[str] = set()

    for thirdparty_pkg in pkg_graph.third_party_dependencies(target_pkg.name):
        assert thirdparty_pkg.name not in thirdparty_pkgs
        depspecs = None
        for p, pc in thirdparty_pkg.invert_dependencies.items():
            if depspecs is None:
                depspecs = pc
            else:
                try:
                    depspecs = find_common_specs(depspecs, pc)
                except IncompatibleDependencyError:
                    logger.error(
                        "Encounter an incompatible dependency {}. Found it in:\n{}",
                        thirdparty_pkg.name,
                        "\n".join(
                            f"\t- {packages[pkgname].location}"
                            for pkgname in thirdparty_pkg.invert_dependencies.keys()
                        ),
                    )

                    if ignore_invalid_dependency:
                        invalid_thirdparty_pkgs.add(thirdparty_pkg.name)
                    else:
                        raise
        assert depspecs is not None
        thirdparty_pkgs[thirdparty_pkg.name] = (
            set(thirdparty_pkg.invert_dependencies.keys()),
            depspecs,
        )

    # now install the target package
    # step 1: gather all dependencies in one file and install it.
    install_pkg_dependencies(
        target_pkg,
        {depname: depspecs for depname, (_, depspecs) in thirdparty_pkgs.items()},
        cfg,
    )

    # step 2: install all local packages in editable mode
    if local_dep_pkgs is None:
        local_dep_pkgs = [pkg for pkg in pkg_graph.local_dependencies(target_pkg.name)]

    pkg_venv_path = venv_path(
        target_pkg.name,
        target_pkg.location,
        cfg.python_virtualenv_path,
        cfg.get_python_path(),
    )

    if len(local_dep_pkgs) > 0:
        logger.info(
            "Installing local packages: {}",
            ", ".join(pkg.name for pkg in local_dep_pkgs),
        )
        for pkg in tqdm(local_dep_pkgs):
            install_bare_pkg(pkg, cfg, pkg_venv_path)

    # step 3: check if we need to build and install any extension module
    # TODO: implement this -- for now, users can use the `build` command to build manually
    logger.info(
        "Finished installing package {} to env {}", target_pkg.name, pkg_venv_path
    )


def install_bare_pkg(pkg: Package, cfg: PBTConfig, virtualenv: Optional[Path] = None):
    """Install a package without any dependencies in editable mode"""
    with mask_file(pkg.location / "pyproject.toml"), mask_file(
        pkg.location / "poetry.lock"
    ):
        with open(pkg.location / "pyproject.toml", "w") as f:
            doc = document()

            tbl = table()
            tbl.add("name", pkg.name)
            tbl.add("version", pkg.version)
            tbl.add("description", "")
            tbl.add("authors", [])
            if sum(int(x != pkg.name) for x in pkg.include) > 0:
                tbl.add("packages", [{"include": x} for x in pkg.include])

            doc.add(SingleKey("tool.poetry", t=KeyType.Bare), tbl)

            tbl = table()
            tbl.add("requires", ["poetry-core>=1.0.0"])
            tbl.add("build-backend", "poetry.core.masonry.api")
            doc.add(nl())
            doc.add("build-system", tbl)

            f.write(dumps(doc))

        install_poetry_package(pkg, cfg, virtualenv, quiet=True)


def install_pkg_dependencies(
    pkg: Package,
    deps: dict[str, DepConstraints],
    cfg: PBTConfig,
    virtualenv: Optional[Path] = None,
):
    with open(cfg.pkg_cache_dir(pkg) / "pyproject.modified.toml", "w") as f:
        doc = document()

        tbl = table()
        tbl.add("name", pkg.name)
        tbl.add("version", pkg.version)
        tbl.add("description", "")
        tbl.add("authors", [])

        if sum(int(x != pkg.name) for x in pkg.include) > 0:
            tbl.add("packages", [{"include": x} for x in pkg.include])

        doc.add(SingleKey("tool.poetry", t=KeyType.Bare), tbl)

        tbl = table()
        if "python" not in deps:
            tbl.add("python", f"=={cfg.get_python_version()}")
        for dep, specs in deps.items():
            tbl.add(dep, serialize_dep_specs(specs))
        doc.add(nl())
        doc.add(SingleKey("tool.poetry.dependencies", t=KeyType.Bare), tbl)

        tbl = table()
        tbl.add("requires", ["poetry-core>=1.0.0"])
        tbl.add("build-backend", "poetry.core.masonry.api")
        doc.add(nl())
        doc.add("build-system", tbl)

        f.write(dumps(doc))

    try:
        os.rename(
            pkg.location / "pyproject.toml",
            cfg.pkg_cache_dir(pkg) / "pyproject.origin.toml",
        )
        if (pkg.location / "poetry.lock").exists():
            os.rename(
                pkg.location / "poetry.lock",
                cfg.pkg_cache_dir(pkg) / "poetry.origin.lock",
            )
        shutil.copy(
            cfg.pkg_cache_dir(pkg) / "pyproject.modified.toml",
            pkg.location / "pyproject.toml",
        )
        if (cfg.pkg_cache_dir(pkg) / "poetry.modified.lock").exists():
            shutil.copy(
                cfg.pkg_cache_dir(pkg) / "poetry.modified.lock",
                pkg.location / "poetry.lock",
            )

        install_poetry_package(pkg, cfg, virtualenv)
    finally:
        os.rename(
            cfg.pkg_cache_dir(pkg) / "pyproject.origin.toml",
            pkg.location / "pyproject.toml",
        )
        if (pkg.location / "poetry.lock").exists():
            os.rename(
                pkg.location / "poetry.lock",
                cfg.pkg_cache_dir(pkg) / "poetry.modified.lock",
            )
        if (cfg.pkg_cache_dir(pkg) / "poetry.origin.lock").exists():
            os.rename(
                cfg.pkg_cache_dir(pkg) / "poetry.origin.lock",
                pkg.location / "poetry.lock",
            )


def install_poetry_package(
    pkg: Package,
    cfg: PBTConfig,
    virtualenv: Optional[Path] = None,
    quiet: bool = False,
):
    if virtualenv is None:
        virtualenv = venv_path(
            pkg.name,
            pkg.location,
            cfg.python_virtualenv_path,
            cfg.get_python_path(),
        )

    env: list[str | NewEnvVar] = [x for x in PASSTHROUGH_ENVS if x != "PATH"]
    for k, v in get_virtualenv_environment_variables(virtualenv).items():
        env.append({"name": k, "value": v})

    if (pkg.location / "poetry.lock").exists():
        try:
            exec(
                "poetry check --lock" + (" -q" if quiet else ""),
                cwd=pkg.location,
                env=env,
            )
        except ExecProcessError:
            logger.debug(
                "Package {} poetry.lock is inconsistent with pyproject.toml, updating lock file...",
                pkg.name,
            )
            exec(
                "poetry lock --no-update" + (" -q" if quiet else ""),
                cwd=pkg.location,
                capture_stdout=False,
                env=env,
            )

    try:
        exec(
            f"poetry install" + (" -q" if quiet else ""),
            cwd=pkg.location,
            capture_stdout=False,
            env=env,
        )
    except ExecProcessError:
        # retry without quiet mode to print out the error.
        exec(
            f"poetry install",
            cwd=pkg.location,
            capture_stdout=False,
            env=env,
        )


def get_virtualenv_environment_variables(virtualenv: Path) -> dict:
    return {
        "VIRTUAL_ENV": str(virtualenv),
        "PATH": str(virtualenv / "bin") + os.pathsep + os.environ.get("PATH", ""),
    }


def find_common_specs(
    depspecs: DepConstraints, another_depspecs: DepConstraints
) -> DepConstraints:
    """Find the common specs between two dependencies."""
    specs = {x.constraint or "": x for x in depspecs}
    anotherspecs = {x.constraint or "": x for x in another_depspecs}

    if not (len(specs) == len(depspecs) and len(anotherspecs) == len(another_depspecs)):
        raise IncompatibleDependencyError(
            f"Two dependencies have duplicated specs: {depspecs} and {another_depspecs}"
        )

    if specs.keys() != anotherspecs.keys():
        raise IncompatibleDependencyError(
            f"Two dependencies have different number of specs: {depspecs} and {another_depspecs}"
        )

    newspecs = []
    for constraint, spec in specs.items():
        anotherspec = anotherspecs[constraint]

        specver = parse_version_spec(spec.version_spec)
        anotherspecver = parse_version_spec(anotherspec.version_spec)

        # we should not have ValueError exception because we have checked the compatibility
        try:
            specver = specver.intersect(anotherspecver)
        except ValueError:
            raise IncompatibleDependencyError(
                f"Two dependencies have incompatible specs: {depspecs} and {another_depspecs}"
            )

        newspecs.append(
            DepConstraint(
                specver.to_pep508_string(),
                constraint,
                spec.version_spec_field,
                spec.origin_spec,
            )
        )
    return newspecs


def serialize_dep_specs(specs: DepConstraints) -> Array:
    items = []
    for spec in specs:
        if spec.origin_spec is None:
            item = spec.version_spec
        else:
            item = inline_table()
            item[cast(str, spec.version_spec_field)] = spec.version_spec
            for k, v in spec.origin_spec.items():
                item[k] = v
        items.append(item)

    if len(items) == 1:
        return items[0]
    return Array(items, Trivia(), multiline=True)


# def load_external_dependencies(pkg: Package, cfg: PBTConfig, packages: dict[str, Package]) -> dict[str, Package]:
