#!/usr/bin/env python3
from __future__ import annotations

import argparse
import os
import re
import subprocess
import sys
from shlex import quote
from typing import TYPE_CHECKING

from packaging.version import InvalidVersion, Version

if TYPE_CHECKING:
    from typing import Iterator

VERSION_FILE = "simpleflow/__init__.py"
MAIN_BRANCH = "main"
CHANGELOG_FILE = "CHANGELOG.md"


def color_msg(color: str, msg: str) -> str:
    colors = {
        "green": "\033[92m",
        "yellow": "\033[93m",
        "red": "\033[91m",
        "blue": "\033[94m",
    }
    if color in colors and sys.stdout.isatty():
        return colors[color] + msg + "\033[0m"
    else:
        return msg


def step(msg: str) -> None:
    print(color_msg("blue", f"* {msg}"))


def fail(message: str) -> None:
    """
    Print a message and exit.
    :param message: message to print
    """
    sys.stderr.write(color_msg("red", f"Error: {message}\nExiting...\n"))
    sys.exit(2)


def execute(command: list[str], ignore: bool = False, log: bool = False, dry_run: bool = False) -> str:
    """
    Execute a command and return the output.
    :param command: command to execute
    :param ignore: whether errors should be ignored (default: False)
    :param log: whether to log commands to stdout (default: False)
    :param dry_run: whether to not execute the command (default: False)
    :return: command output
    """
    if log or dry_run:
        print(f"{'would ' if dry_run else ''}execute: {' '.join(quote(c) for c in command)}")
    if dry_run:
        return ""
    env = os.environ.copy()
    env["LANG"] = "C.UTF-8"
    pr = subprocess.Popen(
        command,
        shell=False,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        env=env,
    )

    (out, error) = pr.communicate()
    if pr.returncode != 0 and not ignore:
        fail(f"Error: executing '{command}', {error.decode('utf-8', 'replace')}")
    return out.decode("utf-8")


def current_branch() -> str:
    """
    Return current branch name.
    """
    branch = execute(["git", "branch", "--show-current"]).strip()
    if branch:
        return branch
    fail("Couldn't find current branch, please don't" " be in 'detached' state when running this.")


def on_main_branch() -> bool:
    """
    Check whether we're on the main branch or not. If you're not on the
    main branch, you're supposed to know what you do!
    """
    return current_branch() == MAIN_BRANCH


def current_version() -> Version:
    with open(VERSION_FILE) as f:
        version_line_regex = re.compile(r"""^__version__\s*=\s*['"](?P<version>.*?)['"]\s*$""")
        for line in f:
            m = version_line_regex.match(line)
            if m:
                return Version(m.group("version"))
    fail(f"Unable to find current version in {VERSION_FILE}")


def increment_version(current: Version) -> Version:
    epoch = current.epoch
    release = current.release
    pre = None
    post = None
    dev = None
    local = current.local

    # Increment least significant part
    if current.dev is not None:
        dev = current.dev + 1
    elif current.post is not None:
        post = current.post + 1
    elif current.pre is not None:
        pre = current.pre[:-1] + (current.pre[-1] + 1,)
    else:
        release = release[:-1] + (release[-1] + 1,)

    parts = []

    # Epoch
    if epoch != 0:
        parts.append(f"{epoch}!")

    # Release segment
    parts.append(".".join(str(x) for x in release))

    # Pre-release
    if pre is not None:
        parts.append("".join(str(x) for x in pre))

    # Post-release
    if post is not None:
        parts.append(f".post{post}")

    # Development release
    if dev is not None:
        parts.append(f".dev{dev}")

    # Local version segment
    if local is not None:
        parts.append(f"+{local}")

    return Version("".join(parts))


def generate_version_file(new_version: Version, dry_run: bool) -> None:
    """
    Generate and modify the simpleflow/__init__.py file.
    """
    with open(VERSION_FILE) as f:
        lines = f.readlines()

    def bump_version_line(line):
        if line.startswith("__version__"):
            return f'__version__ = "{new_version}"\n'
        return line

    lines = [bump_version_line(line) for line in lines]
    joined_lines = "".join(lines)

    if dry_run:
        print(f"Would write:\n{joined_lines}")
    else:
        with open(VERSION_FILE, "w") as f:
            f.write(joined_lines)


def changelog_lines(previous_version: Version) -> Iterator[str]:
    out = execute(
        [
            "git",
            "log",
            "--pretty=format:- %b (%s)",
            "--merges",
            f"{previous_version}..",
        ]
    )
    for line in out.splitlines():
        line = re.sub(r"\(Merge pull request (#\d+)[^)]+\)", r"(\1)", line)
        yield line


def proposed_changelog(previous_version: Version, new_version: Version) -> str:
    return "\n{version}\n{underline}\n\n{content}\n".format(
        version=new_version,
        underline="-" * len(str(new_version)),
        content="\n".join(list(changelog_lines(previous_version))),
    )


def write_changelog(new_content: str, new_version: Version) -> None:
    with open(CHANGELOG_FILE) as f:
        current_changelog = f.readlines()

    # safeguard for not documenting the same tag twice
    tag = str(new_version)
    if tag + "\n" in current_changelog:
        fail(f"The tag {tag} is already present in {CHANGELOG_FILE}")

    # detect where the first sub-title begins, it will be the first version
    # section; we will introduce our new changelog just above
    first_version_line_number = next(idx for idx, line in enumerate(current_changelog) if line.startswith("---")) - 2

    with open(CHANGELOG_FILE, "w") as f:
        for idx, line in enumerate(current_changelog):
            if idx == first_version_line_number:
                f.write(new_content)
            f.write(line)


def generate_changelog(previous_version: Version, new_version: Version, dry_run: bool) -> str:
    proposed = proposed_changelog(previous_version, new_version)
    print(proposed.replace("\n", "\n  "))
    if not dry_run:
        write_changelog(proposed, new_version)
    return proposed


def release_tag(new_version: Version, changes: str, dry_run: bool) -> None:
    """
    Commit and push the branch and tag.
    """
    execute(
        ["git", "commit", "-a", "-m", f"Bump version to {new_version}"],
        log=True,
        dry_run=dry_run,
    )
    annotation_message = f"{new_version}\n\nChangelog:\n{changes}"
    execute(
        ["git", "tag", "-a", str(new_version), "-m", annotation_message],
        log=True,
        dry_run=dry_run,
    )
    execute(["git", "push", "origin", "HEAD"], ignore=True, log=True, dry_run=dry_run)
    execute(
        ["git", "push", "origin", f"{new_version}"],
        ignore=True,
        log=True,
        dry_run=dry_run,
    )


def input_new_version(current: Version) -> Version:
    default_new_version = increment_version(current)
    new_version = None
    while not new_version:
        new_version_str = input(f"New version to release [{default_new_version}]: ")
        if new_version_str:
            try:
                new_version = Version(new_version_str)
            except InvalidVersion as ex:
                print(f"{ex}; Should be PEP 440-compatible (for instance in the form: 1.2.3)")
        else:
            new_version = default_new_version
    return new_version


def main():
    parser = argparse.ArgumentParser(description="Build and upload a new release.")
    parser.add_argument(
        "--version",
        "-V",
        action="version",
        version=f"simpleflow {current_version()}",
        help="display version number",
    )
    parser.add_argument("--dry-run", "-n", action="store_true", help="don't actually do anything")
    group = parser.add_mutually_exclusive_group()
    group.add_argument(
        "--test-pypi",
        "-T",
        action="store_const",
        const="testpypi",
        dest="repository",
        help="upload to TestPyPI",
    )
    group.add_argument(
        "--repository",
        help="repository (package index) to upload the package to",
    )
    group.add_argument(
        "--repository-url",
        help="repository (package index) URL to upload the package to",
    )
    parser.add_argument("--new-version", help="new version number")
    args = parser.parse_args()

    dry_run = args.dry_run

    # check whether on main branch or not
    step("Check current branch")
    if not on_main_branch():
        print("WARNING!")
        print(f"  You're not on the main branch ({MAIN_BRANCH}).")
        print("Please confirm you want to continue [y/N]", end=" ")
        answer = input()
        if not answer.lower().startswith("y"):
            fail("Will not continue as you're not on the main branch")

    step("Detect current/new version")
    current = current_version()
    print(f"Current version: {current}")

    # decide a new version number to release
    if args.new_version:
        new_version = Version(args.new_version)
    else:
        new_version = input_new_version(current)

    # generate new version file
    step(f"Generate version file {VERSION_FILE}")
    generate_version_file(new_version, dry_run)

    # generate changelog
    step(f"Generate {CHANGELOG_FILE}")
    changes = generate_changelog(current, new_version, dry_run)

    # tag version
    step("Release tag")
    release_tag(new_version, changes, dry_run)

    # push package to pypi
    step(f"Generate and push package to {args.repository or args.repository_url or 'pypi'}")
    execute(["python", "setup.py", "bdist_wheel", "sdist"], log=True)
    wheel = f"dist/simpleflow-{new_version}-py3-none-any.whl"
    tar_gz = f"dist/simpleflow-{new_version}.tar.gz"
    cmd = ["twine", "upload", wheel, tar_gz]
    if args.repository:
        cmd += ["--repository", args.repository]
    elif args.repository_url:
        cmd += ["--repository-url", args.repository_url]
    execute(cmd, log=True, dry_run=dry_run)


if __name__ == "__main__":
    main()
