# pyright: strict, reportTypeCommentUsage=false, reportMissingTypeStubs=false
import os
import re
import tempfile

from itertools import chain
from typing import Dict, List, Optional, Set, Tuple, TYPE_CHECKING, cast  # noqa

from urllib.parse import urlparse
from metaflow.debug import debug

# Consider re-adding if resolving fails weirdly.
# from metaflow.exception import retry_exp_backoff

from metaflow._vendor.packaging.requirements import Requirement

from ..env_descr import (
    CondaPackageSpecification,
    EnvType,
    PypiCachePackage,
    PypiPackageSpecification,
    ResolvedEnvironment,
)

from ..pypi_package_builder import PackageToBuild, build_pypi_packages

from ..utils import (
    CondaException,
    arch_id,
    channel_or_url,
    filter_user_reqs_by_markers,
    get_requirement,
    parse_explicit_url_conda,
    parse_explicit_url_pypi,
    pypi_tags_from_arch,
    split_into_dict,
)
from . import Resolver
from .pip_resolver import _FAKE_WHEEL

if TYPE_CHECKING:
    from metaflow_extensions.netflix_ext.plugins.conda.env_descr import (
        PackageSpecification,
    )  # noqa


class CondaLockResolver(Resolver):
    TYPES = ["conda-lock"]

    # @retry_exp_backoff(
    #     (CondaException, OSError),
    #     "Failed to resolve conda-lock environment",
    #     tries=3,
    #     max_delay_per_retry=5,
    # )
    def resolve(
        self,
        env_type: EnvType,
        python_version_requested: str,
        deps: Dict[str, List[str]],
        sources: Dict[str, List[str]],
        extras: Dict[str, List[str]],
        architecture: str,
        builder_envs: Optional[List[ResolvedEnvironment]] = None,
        base_env: Optional[ResolvedEnvironment] = None,
        file_paths: Dict[str, List[str]] = {},
    ) -> Tuple[ResolvedEnvironment, Optional[List[ResolvedEnvironment]]]:
        outfile_name = None
        external_packages = {}
        if base_env:
            local_packages = [
                p for p in base_env.packages if not p.is_downloadable_url()
            ]
            external_packages = {
                p.package_name: p.url
                for p in base_env.packages
                if p.TYPE == "pypi" and p.is_external_url(sources)  # type: ignore[attr-defined]
            }
            if local_packages:
                # We actually only care about things that are not online. Derived packages
                # are OK because we can reconstruct them if needed (or they may even
                # be cached)

                raise CondaException(
                    "Local PYPI packages are not supported in MIXED mode: %s"
                    % ", ".join([p.package_name for p in local_packages])
                )

        # self._start_micromamba_server()

        index_url = self._conda.default_pypi_sources[0]
        # pypi_channels contains everything EXCEPT the pypi default because we don't
        # need to add that to poetry
        pypi_channels = list(
            set(sources.get("pypi", [])).difference(["https://pypi.org/simple"])
        )
        debug.conda_exec("Will add pypi channels: %s" % ", ".join(pypi_channels))

        deps = filter_user_reqs_by_markers(deps, python_version_requested, architecture)
        try:
            # We resolve the environment using conda-lock

            # Write out the TOML file. It's easy enough that we don't use another tool
            # to write it out. We use TOML so that we can disable pypi if needed

            pypi_deps = deps.get("pypi", [])
            conda_deps = list(chain(deps.get("conda", []), deps.get("npconda", [])))

            sys_overrides = split_into_dict(deps.get("sys", []))

            # We only add pip if not present
            if not any([(d == "pip" or d.startswith("pip==")) for d in conda_deps]):
                conda_deps.append("pip")
            # Ditto for uv -- uv 0.7.8 is the last to actually work with python 3.7
            # but this is not listed as a condition so we enforce it here
            # TODO: Remove this once python3.7 is retired
            if not any([(d == "uv" or d.startswith("uv==")) for d in conda_deps]):
                # Duplicate code from filter_user_reqs_by_markers in EnvsResolver.
                # Get the minimum python version supported (major.minor) from the requested specifier.
                spec = get_requirement(python_version_requested).specifier
                candidate_versions = [
                    "3.5",
                    "3.6",
                    "3.7",
                    "3.8",
                    "3.9",
                    "3.10",
                    "3.11",
                    "3.12",
                    "3.13",
                    "3.14",
                ]
                matching_versions = list(spec.filter(candidate_versions))
                if matching_versions:
                    min_supported_py = matching_versions[0]
                else:
                    # fallback to the requested version (eg, ==3.10.5 → 3.10)
                    mm = re.match(r"^.*?([0-9]+)\.([0-9]+)", python_version_requested)
                    if mm:
                        min_supported_py = f"{mm.group(1)}.{mm.group(2)}"
                    else:
                        raise CondaException(
                            "Could not determine minimum supported Python "
                            f"version from {python_version_requested}"
                        )
                if min_supported_py in ["3.5", "3.6", "3.7"]:
                    conda_deps.append("uv==<=0.7.8")
                else:
                    conda_deps.append("uv")

            toml_lines = [
                "[build-system]\n",
                'requires = ["poetry>=0.12"]\n',
                'build-backend = "poetry.masonry.api"\n',
                "\n",
                "[tool.conda-lock]\n",
            ]
            # Add channels
            have_extra_channels = (
                len(
                    set(sources.get("conda", [])).difference(
                        map(channel_or_url, self._conda.default_conda_channels)
                    )
                )
                > 0
            )

            toml_lines.append(
                "channels = [%s]\n"
                % ", ".join(["'%s'" % c for c in sources.get("conda", [])])
            )

            toml_lines.append(
                "pip-repositories = [%s]\n"
                % ", ".join(["'%s'" % c for c in pypi_channels])
            )
            if index_url != "https://pypi.org/simple":
                toml_lines.append("allow-pypi-requests = false\n")

            toml_lines.append("\n")
            # TODO: Maybe we can make this better and only relax if :: is for channels
            # that don't exist in the list
            if any(["::" in d for d in conda_deps]) or have_extra_channels:
                addl_env = {"CONDA_CHANNEL_PRIORITY": "flexible"}
            else:
                addl_env = {}

            # Add deps
            toml_lines.append("[tool.conda-lock.dependencies]\n")
            seen = set()  # type: Set[str]
            for d in conda_deps:
                splits = d.split("==", 1)
                if len(splits) == 2:
                    if splits[0] in seen:
                        raise CondaException(
                            "Duplicate conda dependency %s" % splits[0]
                        )
                    toml_lines.append('"%s" = "%s"\n' % (splits[0], splits[1]))
                    seen.add(splits[0])
                else:
                    if d in seen:
                        raise CondaException("Duplicate conda dependency %s" % d)
                    toml_lines.append('"%s" = "*"\n' % d)
                    seen.add(d)
            toml_lines.append("\n")
            toml_lines.append("[tool.poetry.dependencies]\n")
            # In some cases (when we build packages), we may actually have the same
            # dependency multiple times. We keep just the URL one in this case
            pypi_dep_lines = {}  # type: Dict[str, Dict[str, str]]
            for d in pypi_deps:
                splits = d.split("==", 1)
                # Here we re-parse the requirement. It will be one of the four options:
                #  - <package_name>
                #  - <package_name>[extras]
                #  - <package_name>@<url>
                #  - <package_name>[extras]@<url>

                # We first check if we recorded an external URL for this package --
                # if so, we need to indicate that to conda-lock
                external_url = external_packages.get(splits[0], None)
                if external_url:
                    pypi_dep_lines.setdefault(splits[0], {}).update(
                        {"url": external_url, "url_extras": ""}
                    )
                    continue

                # If not external URL, parse the requirement according to cases above
                parsed_req = Requirement(splits[0])
                if parsed_req.extras:
                    extra_part = "extras = [%s]," % ", ".join(
                        ['"%s"' % e for e in parsed_req.extras]
                    )
                else:
                    extra_part = ""

                version_str = splits[1] if len(splits) == 2 else "*"
                if parsed_req.url:
                    if len(splits) == 2:
                        raise CondaException(
                            "Unexpected version on URL requirement %s" % splits[0]
                        )
                    pypi_dep_lines.setdefault(parsed_req.name, {}).update(
                        {"url": parsed_req.url, "url_extras": extra_part}
                    )
                else:
                    pypi_dep_lines.setdefault(parsed_req.name, {}).update(
                        {"version": version_str, "extras": extra_part}
                    )
            for pypi_name, info in pypi_dep_lines.items():
                if "url" in info:
                    toml_lines.append(
                        '"%s" = {%s = "%s", %s source="pypi"}\n'
                        % (
                            pypi_name,
                            "git" if info["url"].startswith("git+") else "url",
                            info["url"],
                            info["url_extras"],
                        )
                    )
                else:
                    toml_lines.append(
                        '"%s" = {version = "%s", %s source="pypi"}\n'
                        % (pypi_name, info["version"], info["extras"])
                    )

            with tempfile.TemporaryDirectory() as conda_lock_dir:
                outfile_name = "/tmp/conda-lock-gen-%s" % os.path.basename(
                    conda_lock_dir
                )

                args = [
                    "lock",
                    "-f",
                    "pyproject.toml",
                    "-p",
                    architecture,
                    "--filename-template",
                    outfile_name,
                    "-k",
                    "explicit",
                    "--conda",
                ]
                #                if "micromamba_server" in self._conda._bins:
                #                    args.extend([self._bins["micromamba_server"], "--micromamba"])
                #                else:
                conda_exec_type = self._conda.conda_executable_type
                if conda_exec_type:
                    args.append(cast(str, self._conda.binary(conda_exec_type)))

                    if conda_exec_type != "conda":
                        args.append("--%s" % conda_exec_type)
                else:
                    raise CondaException("Could not find conda binary for conda-lock")

                # If arch_id() == architecture, we also use the same virtual packages
                # as the ones that exist on the machine to mimic the current behavior
                # of conda/mamba. If not the same architecture, we pass the system
                # overrides (currently __cuda) down as well.
                if arch_id() == architecture or sys_overrides:
                    lines = ["subdirs:\n", "  %s:\n" % architecture, "    packages:\n"]
                    # Strip out build as it seems to be causing some issues with
                    # some versions of conda-lock when it extracts the GLIBC version.
                    lines.extend(
                        '      %s: "%s"\n' % (virt_pkg, virt_build_str.split("=")[0])
                        for virt_pkg, virt_build_str in self._conda.virtual_packages.items()
                        if virt_pkg not in sys_overrides
                    )
                    lines.extend(
                        '      %s: "%s"\n' % (k, v.split("=")[0])
                        for k, v in sys_overrides.items()
                    )

                    with open(
                        os.path.join(conda_lock_dir, "virtual_yml.spec"),
                        mode="w",
                        encoding="ascii",
                    ) as virtual_yml:
                        virtual_yml.writelines(lines)
                        debug.conda_exec(
                            "Virtual package configuration:\n%s" % "".join(lines)
                        )
                    args.extend(["--virtual-package-spec", "virtual_yml.spec"])

                debug.conda_exec("Build directory: %s" % conda_lock_dir)
                # conda-lock will only consider a `pyproject.toml` as a TOML file which
                # is somewhat annoying.
                with open(
                    os.path.join(conda_lock_dir, "pyproject.toml"),
                    mode="w",
                    encoding="ascii",
                ) as input_toml:
                    input_toml.writelines(toml_lines)
                    debug.conda_exec("TOML configuration:\n%s" % "".join(toml_lines))
                self._conda.call_binary(
                    args, binary="conda-lock", addl_env=addl_env, cwd=conda_lock_dir
                )
            # At this point, we need to read the explicit dependencies in the file created
            emit = False
            packages = []  # type: List["PackageSpecification"]
            packages_to_build = {}  # type: Dict[str, PackageToBuild]
            with open(outfile_name, "r", encoding="utf-8") as out:
                for l in out:
                    if emit:
                        if l.startswith("#"):
                            components = l.split()
                            # Line should be # pip <pkg> @ <url>
                            if len(components) != 5:
                                raise CondaException(
                                    "Unexpected package specification line: %s" % l
                                )
                            debug.conda_exec("Got PYPI package %s" % l)
                            if components[4].startswith("git+"):
                                # This is a GIT URL so we have to build the package
                                # See pip_resolver.py for vcs_info on an explanation
                                # on how we deal with git packages
                                # Looks like gi+https://<url>@<commit>
                                base_build_url = components[4]
                                parse = urlparse(base_build_url)
                                clean_path, clean_commit = parse.path.split("@")
                                clean_url = parse.scheme[4:] + parse.netloc + clean_path
                                base_pkg_url = "%s/%s" % (clean_url, clean_commit)
                                # TODO: Do we need to handle subdirectories
                                cache_base_url = (
                                    PypiCachePackage.make_partial_cache_url(
                                        base_pkg_url, is_real_url=False
                                    )
                                )
                                packages_to_build[cache_base_url] = PackageToBuild(
                                    base_build_url,
                                    PypiPackageSpecification(
                                        _FAKE_WHEEL,
                                        base_pkg_url,
                                        is_real_url=False,
                                        url_format=".whl",
                                    ),
                                )
                                continue
                            # Non-GIT URL
                            parse_result = parse_explicit_url_pypi(components[4])
                            if parse_result.url_format != ".whl":
                                cache_base_url = (
                                    PypiCachePackage.make_partial_cache_url(
                                        parse_result.url, is_real_url=True
                                    )
                                )
                                packages_to_build[cache_base_url] = PackageToBuild(
                                    parse_result.url,
                                    PypiPackageSpecification(
                                        parse_result.filename,
                                        parse_result.url,
                                        is_real_url=True,
                                        url_format=parse_result.url_format,
                                    ),
                                )
                            else:
                                packages.append(
                                    PypiPackageSpecification(
                                        parse_result.filename,
                                        parse_result.url,
                                        url_format=parse_result.url_format,
                                        hashes=(
                                            {parse_result.url_format: parse_result.hash}
                                            if parse_result.hash
                                            else None
                                        ),
                                    )
                                )
                        else:
                            parse_result = parse_explicit_url_conda(l.strip())
                            packages.append(
                                CondaPackageSpecification(
                                    parse_result.filename,
                                    parse_result.url,
                                    url_format=parse_result.url_format,
                                    hashes={
                                        parse_result.url_format: cast(
                                            str, parse_result.hash
                                        )
                                    },
                                )
                            )
                    if not emit and l.strip() == "@EXPLICIT":
                        emit = True
            if packages_to_build:
                with tempfile.TemporaryDirectory() as build_dir:
                    python_version = None  # type: Optional[str]
                    for p in packages:
                        if p.filename.startswith("python-"):
                            python_version = p.package_version
                            break
                    if python_version is None:
                        raise CondaException(
                            "Could not determine version of Python from conda packages"
                        )
                    if architecture == "linux-64":
                        # Get the latest supported GLIBC version so we can generate
                        # the proper tags (only matters on linux)
                        glibc_version_candidates = [
                            d for d in deps.get("sys", []) if d.startswith("__glibc")
                        ]
                        if len(glibc_version_candidates) != 1:
                            raise CondaException(
                                "Could not determine maximum GLIBC version"
                            )
                        # Get version looking like 2.27=0
                        glibc_version = glibc_version_candidates[0][len("__glibc==") :]
                        # Strip =0
                        glibc_version = glibc_version.split("=", 1)[0]
                        # Replace . with _
                        glibc_version = glibc_version.replace(".", "_")
                    else:
                        glibc_version = ""
                    supported_tags = pypi_tags_from_arch(
                        python_version, architecture, glibc_version
                    )
                    if not self._conda.storage:
                        raise CondaException(
                            "Cannot create a relocatable environment as it depends on "
                            "local files or non wheels and no storage backend is defined: %s"
                            % ", ".join([v.url for v in packages_to_build.values()])
                        )
                    built_pypi_packages, builder_envs = build_pypi_packages(
                        self._conda,
                        self._conda.storage,
                        python_version,
                        packages_to_build,
                        builder_envs,
                        build_dir,
                        architecture,
                        supported_tags,
                        sources.get("pypi", []),
                    )
                    packages.extend(built_pypi_packages)
            return (
                ResolvedEnvironment(
                    deps,
                    sources,
                    extras,
                    architecture,
                    all_packages=packages,
                    env_type=env_type,
                ),
                builder_envs,
            )
        finally:
            if outfile_name and os.path.isfile(outfile_name):
                os.unlink(outfile_name)
