import asyncio
import logging
import platform
import re
import subprocess
import typing
from collections import defaultdict
from contextlib import asynccontextmanager
from logging import getLogger
from pathlib import Path
from sys import executable
from tempfile import TemporaryDirectory, gettempdir
from typing import Dict, List, Optional, Set, TextIO, Union
from urllib.parse import urlparse

from filelock import BaseFileLock, FileLock
from packaging.utils import parse_wheel_filename
from rich.progress import Progress

from coiled._beta.core import CloudBeta
from coiled._beta.widgets.util import simple_progress
from coiled.scan import PackageInfo, scan_prefix
from coiled.types import PackageLevelEnum, ResolvedPackageInfo
from coiled.utils import validate_wheel

logger = getLogger("coiled.package_sync")
subdir_datas = {}
PYTHON_VERSION = platform.python_version_tuple()
ANY_AVAILABLE = "ANY-AVAILABLE"


async def create_subprocess_shell(
    cmd: str,
    stdout: Union[TextIO, int, None] = None,
    stderr: Union[TextIO, int, None] = None,
) -> subprocess.CompletedProcess:
    # asyncio.create_subprocess_shell is
    # not supported on windows
    # and broken on py3.7
    loop = asyncio.get_running_loop()
    result = loop.run_in_executor(None, lambda: subprocess.run(cmd, shell=True, stdout=stdout, stderr=stderr))
    return await result


class PackageBuildError(Exception):
    pass


WHEEL_BUILD_LOCKS: Dict[str, BaseFileLock] = {}


async def default_python() -> PackageInfo:
    python_version = platform.python_version()
    return {
        "name": "python",
        "path": None,
        "source": "conda",
        "channel_url": ANY_AVAILABLE,
        "channel": ANY_AVAILABLE,
        "subdir": "linux-64",
        "conda_name": "python",
        "version": python_version,
        "wheel_target": None,
    }


@asynccontextmanager
async def async_lock(lock: BaseFileLock):
    loop = asyncio.get_event_loop()
    ctx = await loop.run_in_executor(None, lock.acquire)
    with ctx:
        yield


async def create_wheel(pkg_name: str, version: str, src: str) -> ResolvedPackageInfo:
    lock = WHEEL_BUILD_LOCKS.setdefault(pkg_name, FileLock(Path(gettempdir()) / (pkg_name + version + ".lock")))
    async with async_lock(lock):
        tmpdir = TemporaryDirectory()
        outdir = Path(tmpdir.name) / Path(pkg_name)
        logger.info(f"Attempting to create a wheel for {pkg_name} @ {src}")
        # must use executable to avoid using some other random python
        p = await create_subprocess_shell(
            cmd=f"{executable} -m pip wheel --wheel-dir {outdir} --no-deps --use-pep517 --no-cache-dir {src}",
            stderr=subprocess.STDOUT,
            stdout=subprocess.PIPE,
        )
        if p.returncode:
            print(f"---Wheel Build Log for {pkg_name}---\n{p.stdout.decode()}")
            return {
                "name": pkg_name,
                "source": "pip",
                "channel": None,
                "conda_name": None,
                "client_version": version,
                "specifier": "",
                "include": False,
                "error": (
                    "Failed to build a wheel for the"
                    " package, will not be included in environment, check stdout for the build log"
                ),
                "note": None,
                "sdist": None,
                "md5": None,
            }
        wheel_fn = next(file for file in outdir.iterdir() if file.suffix == ".whl")
        _, build_version, _, _ = parse_wheel_filename(str(wheel_fn.name))
        has_python, md5 = await validate_wheel(wheel_fn)
    return {
        "name": pkg_name,
        "source": "pip",
        "channel": None,
        "conda_name": None,
        "client_version": str(build_version),
        "specifier": "",
        "include": True,
        "error": None if has_python else "Built wheel contains no python files!",
        "note": f"Wheel built from {src}",
        "sdist": wheel_fn.open("rb"),
        "md5": md5,
    }


async def create_wheel_from_egg(pkg_name: str, version: str, src: str) -> ResolvedPackageInfo:
    tmpdir = TemporaryDirectory()
    outdir = Path(tmpdir.name) / Path(pkg_name)
    outdir.mkdir(parents=True)
    logger.info(f"Attempting to create a wheel for {pkg_name} in directory {src}")
    # must use executable to avoid using some other random python
    p = await create_subprocess_shell(
        cmd=f"{executable} -m wheel convert --dest-dir {outdir} {src}",
        stderr=subprocess.STDOUT,
        stdout=subprocess.PIPE,
    )
    if p.returncode:
        print(f"---Egg to wheel conversion Log for {pkg_name}---\n{p.stdout.decode()}")
        return {
            "name": pkg_name,
            "source": "pip",
            "channel": None,
            "conda_name": None,
            "client_version": version,
            "specifier": "",
            "include": False,
            "error": (
                "Failed to convert the package egg to a wheel"
                ", will not be included in environment, check stdout for egg conversion log"
            ),
            "note": None,
            "sdist": None,
            "md5": None,
        }
    wheel_fn = next(file for file in outdir.iterdir() if file.suffix == ".whl")
    has_python, md5 = await validate_wheel(Path(wheel_fn))
    return {
        "name": pkg_name,
        "source": "pip",
        "channel": None,
        "conda_name": None,
        "client_version": version,
        "specifier": "",
        "include": True,
        "error": None if has_python else "Built wheel has no python files!",
        "note": "Wheel built from local egg",
        "sdist": wheel_fn.open("rb"),
        "md5": md5,
    }


async def approximate_packages(
    cloud: CloudBeta,
    packages: List[PackageInfo],
    priorities: Dict[str, PackageLevelEnum],
    progress: Optional[Progress] = None,
    strict: bool = False,
) -> typing.List[ResolvedPackageInfo]:
    user_conda_installed_python = next((p for p in packages if p["name"] == "python"), None)
    user_conda_installed_pip = next(
        (i for i, p in enumerate(packages) if p["name"] == "pip" and p["source"] == "conda"),
        None,
    )
    if not user_conda_installed_pip:
        # This means pip was installed by pip, or the system
        # package manager
        # Insert a conda version of pip to be installed first, it will
        # then be used to install the users version of pip
        pip = next(
            (p for p in packages if p["name"] == "pip" and p["source"] == "pip"),
            None,
        )
        if not pip:
            # insert a modern version and hope it does not introduce conflicts
            packages.append(
                {
                    "name": "pip",
                    "path": None,
                    "source": "conda",
                    "channel_url": "https://conda.anaconda.org/conda-forge/",
                    "channel": "conda-forge",
                    "subdir": "noarch",
                    "conda_name": "pip",
                    "version": "22.3.1",
                    "wheel_target": None,
                }
            )
        else:
            # insert the users pip version and hope it exists on conda-forge
            packages.append(
                {
                    "name": "pip",
                    "path": None,
                    "source": "conda",
                    "channel_url": "https://conda.anaconda.org/conda-forge/",
                    "channel": "conda-forge",
                    "subdir": "noarch",
                    "conda_name": "pip",
                    "version": pip["version"],
                    "wheel_target": None,
                }
            )
    coiled_selected_python = None
    if not user_conda_installed_python:
        # insert a special python package
        # that the backend will pick a channel for
        coiled_selected_python = await default_python()
        packages.append(coiled_selected_python)
    with simple_progress("Validating environment", progress=progress):
        results = await cloud._approximate_packages(
            packages=[
                {
                    "name": pkg["name"],
                    "priority_override": 100 if (strict or pkg["wheel_target"]) else priorities.get(pkg["name"]),
                    "python_major_version": PYTHON_VERSION[0],
                    "python_minor_version": PYTHON_VERSION[1],
                    "python_patch_version": PYTHON_VERSION[2],
                    "source": pkg["source"],
                    "channel_url": pkg["channel_url"],
                    "channel": pkg["channel"],
                    "subdir": pkg["subdir"],
                    "conda_name": pkg["conda_name"],
                    "version": pkg["version"],
                    "wheel_target": pkg["wheel_target"],
                }
                for pkg in packages
            ]
        )
    result_map = {r["name"]: r for r in results}
    finalized_packages: typing.List[ResolvedPackageInfo] = []

    if not user_conda_installed_python and coiled_selected_python:
        # user has no python version installed by conda
        # we should have a result of asking the backend to
        # pick conda channel that has the users python version
        if result_map["python"]["error"]:
            finalized_packages.append(
                {
                    "name": "python",
                    "source": "conda",
                    "channel": None,
                    "conda_name": "python",
                    "client_version": coiled_selected_python["version"],
                    "specifier": result_map["python"]["specifier"] or "",
                    "include": result_map["python"]["include"],
                    "note": None,
                    "error": result_map["python"]["error"],
                    "sdist": None,
                    "md5": None,
                }
            )
        else:
            assert result_map["python"]["note"]
            channel_url, channel = result_map["python"]["note"].split(",")
            finalized_packages.append(
                {
                    "name": "python",
                    "source": "conda",
                    "channel": channel,
                    "conda_name": "python",
                    "client_version": coiled_selected_python["version"],
                    "specifier": result_map["python"]["specifier"] or "",
                    "include": result_map["python"]["include"],
                    "note": None,
                    "error": result_map["python"]["error"],
                    "sdist": None,
                    "md5": None,
                }
            )
        # we can pull our special python package out of the list
        # now
        packages.remove(coiled_selected_python)
    for pkg in packages:

        if pkg["wheel_target"] and result_map[pkg["name"]]["include"]:
            p = urlparse(pkg["wheel_target"])
            if p.scheme == "" and not Path(pkg["wheel_target"]).exists():
                # lack of scheme means a local file
                # sometimes the wheel target taken from
                # direct_url.json references a file that does not exist
                # skip over trying to sync that and treat it like a normal package
                # this can happen in conda/system packages
                # where the package metadata was generated on a build system and not locally
                finalized_packages.append(
                    {
                        "name": pkg["name"],
                        "source": pkg["source"],
                        "channel": pkg["channel"],
                        "conda_name": pkg["conda_name"],
                        "client_version": pkg["version"],
                        "specifier": result_map[pkg["name"]]["specifier"] or "",
                        "include": result_map[pkg["name"]]["include"],
                        "note": result_map[pkg["name"]]["note"],
                        "error": result_map[pkg["name"]]["error"],
                        "sdist": None,
                        "md5": None,
                    }
                )
            elif pkg["wheel_target"].endswith(".egg"):
                with simple_progress(f'Creating wheel from egg for {pkg["name"]}', progress=progress):
                    finalized_packages.append(
                        await create_wheel_from_egg(
                            pkg_name=pkg["name"],
                            version=pkg["version"],
                            src=pkg["wheel_target"],
                        )
                    )
            else:
                with simple_progress(f'Creating wheel for {pkg["name"]}', progress=progress):
                    finalized_packages.append(
                        await create_wheel(
                            pkg_name=pkg["name"],
                            version=pkg["version"],
                            src=pkg["wheel_target"],
                        )
                    )
        else:
            finalized_packages.append(
                {
                    "name": pkg["name"],
                    "source": pkg["source"],
                    "channel": pkg["channel"],
                    "conda_name": pkg["conda_name"],
                    "client_version": pkg["version"],
                    "specifier": result_map[pkg["name"]]["specifier"] or "",
                    "include": result_map[pkg["name"]]["include"],
                    "note": result_map[pkg["name"]]["note"],
                    "error": result_map[pkg["name"]]["error"],
                    "sdist": None,
                    "md5": None,
                }
            )
    return finalized_packages


pip_bad_req_regex = (
    r"(?P<package>.+) (?P<version>.+) has requirement "
    r"(?P<requirement>.+), but you have (?P<requirement2>.+) (?P<reqversion>.+)."
)


async def check_pip_happy(progress: Optional[Progress] = None) -> Dict[str, List[str]]:
    with simple_progress("Running pip check", progress=progress):
        proc = await create_subprocess_shell(
            f"{executable} -m pip check", stdout=subprocess.PIPE, stderr=subprocess.PIPE
        )
        faulty_packages = defaultdict(list)
        if proc.returncode:
            output = proc.stdout.decode()
            bad_reqs = re.finditer(pip_bad_req_regex, output)
            for bad_req in bad_reqs:
                groups = bad_req.groupdict()
                span = bad_req.span()
                warning = output[span[0] : span[1]]
                logger.warning(warning)
                faulty_packages[groups["package"]].append(warning)
        return faulty_packages


async def create_environment_approximation(
    cloud: CloudBeta,
    priorities: Dict[str, PackageLevelEnum],
    only: Optional[Set[str]] = None,
    strict: bool = False,
    progress: Optional[Progress] = None,
) -> typing.List[ResolvedPackageInfo]:
    packages = await scan_prefix(progress=progress)
    bad_packages = await check_pip_happy(progress)
    if bad_packages:
        packages = [pkg for pkg in packages if pkg["name"] not in bad_packages]
    if only:
        packages = filter(lambda pkg: pkg["name"] in only, packages)
    # TODO: private conda channels
    result = await approximate_packages(
        cloud=cloud, packages=[pkg for pkg in packages], priorities=priorities, strict=strict, progress=progress
    )
    for bad_package, error_list in bad_packages.items():
        errors = "\n".join(error_list)
        result.append(
            {
                "name": bad_package,
                "error": f"Pip check had the following issues that need resolving: \n{errors}",
                "source": "pip",
                "channel": None,
                "client_version": "",
                "conda_name": None,
                "include": False,
                "md5": None,
                "note": None,
                "sdist": None,
                "specifier": "",
            }
        )
    return result


if __name__ == "__main__":
    from logging import basicConfig

    basicConfig(level=logging.INFO)

    from rich.console import Console
    from rich.table import Table

    async def run():
        async with CloudBeta(asynchronous=True) as cloud:
            return await create_environment_approximation(
                cloud=cloud,
                priorities={
                    "dask": PackageLevelEnum.CRITICAL,
                    "twisted": PackageLevelEnum.IGNORE,
                    "graphviz": PackageLevelEnum.OK,
                    "icu": PackageLevelEnum.OK,
                },
            )

    result = asyncio.run(run())

    table = Table(title="Packages")
    keys = ("name", "source", "include", "client_version", "specifier", "error", "note")

    for key in keys:
        table.add_column(key)

    for pkg in result:
        row_values = [str(pkg.get(key, "")) for key in keys]
        table.add_row(*row_values)
    console = Console()
    console.print(table)
    console.print(table)
