#!/usr/bin/env python3

from __future__ import absolute_import, annotations

import atexit
import functools
import glob
import hashlib
import io
import os
import shutil
import subprocess
import sys
import tempfile
from argparse import Action, ArgumentDefaultsHelpFormatter, ArgumentError, ArgumentParser
from dataclasses import dataclass
from email.parser import Parser
from enum import Enum, unique
from http.server import HTTPServer, SimpleHTTPRequestHandler
from pathlib import Path, PurePath
from typing import Dict, Iterator, Optional, Tuple

from package.scie_config import PlatformConfig, ScieConfig
from pex.common import safe_mkdtemp
from pex.jobs import Job, SpawnedJob, execute_parallel
from pex.pip.version import PipVersion
from pex.sysconfig import SysPlatform

DEFAULT_DIST_DIR = Path("dist")
PACKAGE_DIR = Path("package")


def build_pex_pex(
    output_file: PurePath, verbosity: int = 0, env: Optional[Dict[str, str]] = None
) -> PurePath:
    # NB: We do not include the subprocess extra (which would be spelled: `.[subprocess]`) since we
    # would then produce a pex that would not be consumable by all python interpreters otherwise
    # meeting `python_requires`; ie: we'd need to then come up with a deploy environment / deploy
    # tooling, that built subprocess32 for linux cp27m, cp27mu, pypy, ... etc. Even with all the
    # work expended to do this, we'd still miss some platform someone wanted to run the Pex PEX on.
    # As such, we just ship unadorned Pex which is pure-python and universal. Any user wanting the
    # extra is encouraged to build a Pex PEX for their particular platform themselves.
    pex_requirement = "."

    args = [
        sys.executable,
        "-m",
        "pex",
        *["-v" for _ in range(verbosity)],
        "--disable-cache",
        "--no-build",
        "--no-compile",
        "--no-use-system-time",
        "--python-shebang",
        "/usr/bin/env python",
        "--no-strip-pex-env",
        "--include-tools",
        "--record-git-state",
        "--pip-version",
        "latest-compatible",
        "-o",
        str(output_file),
        "-c",
        "pex",
        pex_requirement,
    ]
    subprocess.run(args=args, env=env, check=True)
    return output_file


def build_pex_scie(
    platform_info: tuple[PlatformConfig, Path],
    *,
    lock: Path,
    pex_requirement: str,
    verbosity: int = 0,
    env: Optional[Dict[str, str]] = None,
) -> SpawnedJob[tuple[str, PlatformConfig]]:
    platform_config, complete_platform = platform_info
    command = [
        sys.executable,
        "-m",
        "pex",
        *["-v" for _ in range(verbosity)],
        "--disable-cache",
        "--no-build",
        "--no-compile",
        "--no-use-system-time",
        "--record-git-state",
        "--venv",
        "--no-strip-pex-env",
        "--complete-platform",
        str(complete_platform),
        "--lock",
        str(lock),
        "--scie",
        "eager",
        "--scie-only",
        "--scie-name-style",
        "platform-file-suffix",
        "--scie-platform",
        platform_config.name,
        "--scie-pbs-release",
        platform_config.pbs_release,
        "--scie-python-version",
        platform_config.python_version,
        "--scie-pbs-stripped",
        "--scie-hash-alg",
        "sha256",
        "--scie-busybox",
        "@pex",
        "--scie-busybox-pex-entrypoint-env-passthrough",
        "--pip-version",
        PipVersion.LATEST_COMPATIBLE.value,
        "-o",
        os.path.join(safe_mkdtemp(), "pex"),
        "-c",
        "pex",
        "--project",
        pex_requirement,
    ]
    return SpawnedJob.stdout(
        job=Job(
            command=command, process=subprocess.Popen(args=command, env=env, stdout=subprocess.PIPE)
        ),
        result_func=lambda stdout: (stdout.decode("utf-8").strip(), platform_config),
    )


def build_pex_scies(
    dist_dir: Path,
    scie_choice: ScieChoice,
    verbosity: int = 0,
    env: Optional[Dict[str, str]] = None,
) -> Iterator[tuple[Path, PlatformConfig]]:
    scie_config = ScieConfig.load()

    pex_requirement = f".[{','.join(scie_config.pex_extras)}]"

    lock = PACKAGE_DIR / "pex-scie.lock"
    if not lock.exists():
        raise SystemExit(
            f"The Pex scie lock at {lock} does not exist.\n"
            f"Run `uv run dev-cmd gen-scie-platform -- --all ...` to generate it."
        )

    missing_platforms: list[str] = []
    platforms: list[tuple[PlatformConfig, Path]] = []
    for platform_config in scie_config.platforms:
        if (
            scie_choice is ScieChoice.CURRENT
            and platform_config.platform is not SysPlatform.CURRENT
        ):
            continue

        complete_platform = PACKAGE_DIR / "complete-platforms" / f"{platform_config.name}.json"
        if complete_platform.exists():
            platforms.append((platform_config, complete_platform))
        elif platform_config.required:
            missing_platforms.append(platform_config.name)

    if missing_platforms:
        missing = "\n".join(
            f"{index}. {missing_platform}"
            for index, missing_platform in enumerate(missing_platforms, start=1)
        )
        raise SystemExit(
            f"Of the {len(scie_config.platforms)} expected Pex scie complete platforms, "
            f"{len(missing_platforms)} {'is' if len(missing_platforms) == 1 else 'are'} missing:\n"
            f"{missing}"
        )

    for output_file, platform_config in execute_parallel(
        platforms,
        spawn_func=functools.partial(
            build_pex_scie, verbosity=verbosity, lock=lock, pex_requirement=pex_requirement, env=env
        ),
    ):
        artifacts = glob.glob(f"{output_file}*")
        scie_artifacts = [artifact for artifact in artifacts if not artifact.endswith(".sha256")]
        if len(scie_artifacts) != 1:
            artifact_lines = "\n".join(sorted(artifacts))
            raise SystemExit(
                f"Found unexpected artifacts after generating Pex scie:\n{artifact_lines}"
            )
        scie_name = os.path.basename(scie_artifacts[0])
        for artifact in artifacts:
            shutil.move(artifact, dist_dir / os.path.basename(artifact))

        yield dist_dir / scie_name, platform_config


def describe_rev() -> str:
    if not os.path.isdir(".git") and os.path.isfile("PKG-INFO"):
        # We're being build from an unpacked sdist.
        with open("PKG-INFO") as fp:
            return Parser().parse(fp).get("Version", "Unknown Version")

    git_describe = subprocess.run(
        ["git", "describe"], check=True, stdout=subprocess.PIPE, encoding="utf-8"
    )
    return git_describe.stdout.strip()


def describe_file(path: Path) -> Tuple[str, int]:
    hasher = hashlib.sha256()
    size = 0
    with path.open("rb") as fp:
        for chunk in iter(lambda: fp.read(io.DEFAULT_BUFFER_SIZE), b""):
            hasher.update(chunk)
            size += len(chunk)

    return hasher.hexdigest(), size


@dataclass
class _FormatData:
    _name: str
    build_args: Tuple[str, ...]
    build_env: Tuple[Tuple[str, str], ...]


@unique
class Format(_FormatData, Enum):
    @classmethod
    def for_name(cls, name: str):
        for member in cls:
            if name in (member._name, member.name):
                return member
        raise ValueError(f"Not a recognized format: {name}")

    SDIST = ("sdist", ("--sdist",), ())
    WHEEL = (
        "whl",
        (
            "--wheel",
            "--config-setting=--build-option=--python-tag=py2.py35.py36.py37.py38.py39.py310.py311",
        ),
        (),
    )
    WHEEL_3_12_PLUS = (
        "whl-3.12-plus",
        ("--wheel", "--config-setting=--build-option=--python-tag=py3.py312"),
        (("__PEX_BUILD_WHL_3_12_PLUS__", "1"),),
    )

    def add_build_env(self, env: Optional[Dict[str, str]] = None) -> Optional[Dict[str, str]]:
        if not env and not self.build_env:
            return None
        build_env = dict(env or os.environ)
        build_env.update(self.build_env)
        return build_env

    def __str__(self) -> str:
        return self._name


def build_pex_dists(
    dist_dir: Path,
    dist_fmt: Format,
    *additional_dist_fmts: Format,
    verbose: bool = False,
    env: Optional[Dict[str, str]] = None
) -> Iterator[PurePath]:
    tmp_dir = dist_dir / ".tmp"
    tmp_dir.mkdir(parents=True, exist_ok=True)
    atexit.register(shutil.rmtree, tmp_dir, ignore_errors=True)
    out_dir = tempfile.mkdtemp(dir=tmp_dir)

    output = None if verbose else subprocess.DEVNULL

    for fmt in [dist_fmt, *additional_dist_fmts]:
        subprocess.run(
            args=[sys.executable, "-m", "build", "--outdir", out_dir, *fmt.build_args],
            env=fmt.add_build_env(env),
            stdout=output,
            stderr=output,
            check=True,
        )

    for dist in os.listdir(out_dir):
        built = dist_dir / dist
        shutil.move(os.path.join(out_dir, dist), built)
        yield built


def main(
    dist_dir: Path,
    *additional_dist_formats: Format,
    verbosity: int = 0,
    embed_docs: bool = False,
    clean_docs: bool = False,
    pex_output_file: Optional[Path] = None,
    scie_choice: ScieChoice | None = None,
    markdown_hash_table_file: Optional[Path] = None,
    serve: bool = False
) -> None:
    env = os.environ.copy()
    if embed_docs:
        env.update(__PEX_BUILD_INCLUDE_DOCS__="1")

    hash_table: dict[Path, tuple[str, int]] = {}
    if pex_output_file:
        print(f"Building Pex PEX to `{pex_output_file}` ...")
        build_pex_pex(pex_output_file, verbosity, env=env)

        rev = describe_rev()
        sha256, size = describe_file(pex_output_file)
        with (pex_output_file.parent / f"{pex_output_file.name}.sha256").open("w") as fp:
            fp.write(f"{sha256} *{pex_output_file.name}")
        hash_table[pex_output_file] = sha256, size
        print(f"Built Pex PEX @ {rev}:")
        print(f"sha256: {sha256}")
        print(f"  size: {size}")

    if scie_choice:
        scies = "scies" if scie_choice is ScieChoice.ALL else "scie"
        print(f"Building Pex {scies} to `{dist_dir}` ...")
        for scie, platform in build_pex_scies(dist_dir, scie_choice, verbosity, env=env):
            hash_table[scie] = describe_file(scie)
            print(f"  Built Pex scie for {platform.name} at `{scie}`")

    if markdown_hash_table_file and hash_table:
        with markdown_hash_table_file.open(mode="w") as fp:
            print("|file|sha256|size|", file=fp)
            print("|----|------|----|", file=fp)
            for file, (sha256, size) in sorted(hash_table.items()):
                print(f"|{file.name}|{sha256}|{size}|", file=fp)

        print(f"Generated markdown table of Pex sizes & hashes at `{markdown_hash_table_file}`")

    if additional_dist_formats:
        print(
            f"Building additional distribution formats to `{dist_dir}`: "
            f'{", ".join(f"{i + 1}.) {fmt}" for i, fmt in enumerate(additional_dist_formats))} ...'
        )
        built = list(
            build_pex_dists(
                dist_dir,
                additional_dist_formats[0],
                *additional_dist_formats[1:],
                verbose=verbosity > 0,
                env=env
            )
        )
        print("Built:")
        for dist_path in built:
            print(f"  {dist_path}")

    if clean_docs:
        shutil.rmtree(dist_dir / "docs", ignore_errors=True)

    if serve:
        server = HTTPServer(("", 0), SimpleHTTPRequestHandler)
        host, port = server.server_address

        print(f"Serving Pex distributions from `{dist_dir}` at http://{host}:{port} ...")

        os.chdir(dist_dir)
        try:
            server.serve_forever()
        except KeyboardInterrupt:
            print(f"Server shut down in response to keyboard interrupt.")


class ScieChoice(Enum):
    ALL = "all"
    CURRENT = "current"


class HandleScies(Action):
    def __init__(self, *args, **kwargs):
        kwargs["nargs"] = 0
        super().__init__(*args, **kwargs)

    def __call__(self, parser, namespace, values, option_string=None):
        scie_choice = None
        if option_string == "--scie":
            scie_choice = ScieChoice.CURRENT
        elif option_string == "--scies":
            scie_choice = ScieChoice.ALL
        elif option_string:
            raise ArgumentError(self, f"Unknown option {option_string}")
        setattr(namespace, self.dest, scie_choice)


if __name__ == "__main__":
    parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
    parser.add_argument(
        "-v", dest="verbosity", action="count", default=0, help="Increase output verbosity level."
    )
    parser.add_argument(
        "--embed-docs",
        default=False,
        action="store_true",
        help="Embed offline docs in the built binary distributions.",
    )
    parser.add_argument(
        "--clean-docs",
        default=False,
        action="store_true",
        help="Clean up loose generated docs after they have embedded in binary distributions.",
    )
    parser.add_argument(
        "--additional-format",
        dest="additional_formats",
        choices=list(Format),
        type=Format.for_name,
        action="append",
        help="Package Pex in additional formats.",
    )
    parser.add_argument(
        "--no-pex",
        default=False,
        action="store_true",
        help="Do not build the Pex PEX.",
    )
    parser.add_argument(
        "-d",
        "--dist-dir",
        default=DEFAULT_DIST_DIR,
        type=Path,
        help="Build packages to this directory.",
    )
    parser.add_argument(
        "--pex-output-file",
        default=None,
        type=Path,
        help="Build the Pex PEX at this path (default: --dist-dir / 'pex').",
    )
    parser.add_argument(
        "--scie",
        "--scies",
        dest="scies",
        default=None,
        action=HandleScies,
        help=(
            "Build PEX scies. Passing --scie builds the PEX scie for the current platform and "
            "passing --scies builds PEX scies for all supported platforms."
        ),
    )
    parser.add_argument(
        "--gen-md-table-of-hash-and-size",
        default=None,
        type=Path,
        help="A path to generate a markdown table of packaged asset hashes to.",
    )
    parser.add_argument(
        "--serve",
        default=False,
        action="store_true",
        help="After packaging Pex serve up the packages over HTTP.",
    )
    args = parser.parse_args()

    main(
        args.dist_dir,
        *(args.additional_formats or ()),
        verbosity=args.verbosity,
        embed_docs=args.embed_docs,
        clean_docs=args.clean_docs,
        pex_output_file=None if args.no_pex else (args.pex_output_file or args.dist_dir / "pex"),
        scie_choice=args.scies,
        markdown_hash_table_file=args.gen_md_table_of_hash_and_size,
        serve=args.serve
    )
