# pyright: strict, reportTypeCommentUsage=false, reportMissingTypeStubs=false

import os
import platform

from shutil import which
from typing import NamedTuple, Optional, Sequence, Tuple, Union
from urllib.parse import urlparse

from metaflow.exception import MetaflowException
import metaflow.metaflow_config as mf_config
from metaflow.metaflow_config import (
    CONDA_MAGIC_FILE_V2,  # type: ignore
    CONDA_PREFERRED_FORMAT,  # type: ignore
)

from metaflow.metaflow_environment import InvalidEnvironmentException

# NOTA: Most of the code does not assume that there are only two formats BUT the
# transmute code does (since you can only specify the infile -- the outformat and file
# are inferred)
_ALL_CONDA_FORMATS = (".tar.bz2", ".conda")

# List of formats that guarantees the preferred format is first. This is important as
# functions that rely on selecting the "preferred" source of a package rely on the
# preferred format being first.
if CONDA_PREFERRED_FORMAT:
    CONDA_FORMATS = (
        CONDA_PREFERRED_FORMAT,
        *[x for x in _ALL_CONDA_FORMATS if x != CONDA_PREFERRED_FORMAT],
    )  # type: Tuple[str, ...]
else:
    CONDA_FORMATS = CONDA_PREFERRED_FORMAT  # type: Tuple[str, ...]
TRANSMUT_PATHCOMPONENT = "_transmut"


# On Linux systems, called md5sum and md5 on mac
_md5sum_path = which("md5sum") or which("md5")


class CondaException(MetaflowException):
    headline = "Conda ran into an error while setting up environment."

    def __init__(self, error: Union[Sequence[Exception], str]):
        if isinstance(error, list):
            error = "\n".join([str(x) for x in error])
        super(CondaException, self).__init__(error)  # type: ignore


class CondaStepException(CondaException):
    def __init__(self, exception: CondaException, steps: Sequence[str]):
        msg = "Step(s): {steps}, Error: {error}".format(
            steps=steps, error=exception.message
        )
        super(CondaStepException, self).__init__(msg)


def convert_filepath(path: str, file_format: Optional[str] = None) -> Tuple[str, str]:
    if file_format and not path.endswith(file_format):
        for f in CONDA_FORMATS:
            if path.endswith(f):
                path = path[: -len(f)] + file_format
                break
        else:
            raise ValueError(
                "URL '%s' does not end with a supported file format %s"
                % (path, str(CONDA_FORMATS))
            )
    return os.path.split(path)


def get_conda_manifest_path(ds_root: str) -> str:
    return os.path.join(ds_root, CONDA_MAGIC_FILE_V2)  # type: ignore


def get_conda_root(datastore_type: str) -> str:
    conda_root = getattr(mf_config, "CONFIG_%sROOT" % datastore_type.upper())
    if conda_root is None:
        # We error on METAFLOW_DATASTORE_SYSROOT_<ds> because that is the default used
        raise MetaflowException(
            msg="METAFLOW_DATASTORE_SYSROOT_%s must be set!" % datastore_type.upper()
        )
    return conda_root  # type: ignore


def arch_id() -> str:
    bit = "32"
    if platform.machine().endswith("64"):
        bit = "64"
    if platform.system() == "Linux":
        return "linux-%s" % bit
    elif platform.system() == "Darwin":
        # Support M1 Macmetaf
        if platform.machine() == "arm64":
            return "osx-arm64"
        else:
            return "osx-%s" % bit
    else:
        raise InvalidEnvironmentException(
            "The *@conda* decorator is not supported "
            "outside of Linux and Darwin platforms"
        )


ParseExplicitResult = NamedTuple(
    "ParseExplicitResult",
    [("filename", str), ("url", str), ("url_format", str), ("hash", str)],
)


def parse_explicit_url_conda(url: str) -> ParseExplicitResult:
    # Takes a URL in the form url#hash and returns:
    #  - the filename
    #  - the URL (without the hash)
    #  - the format for the URL
    #  - the hash
    filename = None
    url_format = None

    url_clean, url_hash = url.rsplit("#", 1)

    filename = os.path.split(urlparse(url_clean).path)[1]
    for f in _ALL_CONDA_FORMATS:
        if filename.endswith(f):
            url_format = f
            filename = filename[: -len(f)]
            break
    else:
        raise CondaException(
            "URL '%s' is not a supported format (%s)" % (url, CONDA_FORMATS)
        )
    return ParseExplicitResult(
        filename=filename, url=url_clean, url_format=url_format, hash=url_hash
    )


def parse_explicit_url_pip(url: str) -> ParseExplicitResult:
    # Takes a URL in the form url#hash and returns:
    #  - the filename
    #  - the URL (without the hash)
    #  - the format for the URL
    #  - the hash
    filename = None
    url_format = None

    url_clean, url_hash = url.rsplit("#", 1)
    if not url_hash.startswith("sha256="):
        raise CondaException("URL '%s' has a SHA type which is not supported" % url)
    url_hash = url_hash[7:]

    filename = os.path.split(urlparse(url_clean).path)[1]
    for f in [".whl", ".tar.gz"]:
        if filename.endswith(f):
            url_format = f
            filename = filename[: -len(f)]
            break
    else:
        raise CondaException(
            "URL '%s' is not a supported format (%s)" % (url, str([".whl", ".tar.gz"]))
        )
    return ParseExplicitResult(
        filename=filename, url=url_clean, url_format=url_format, hash=url_hash
    )


def plural_marker(count: int) -> str:
    return "s" if count != 1 else ""
