# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.


import argparse
import logging
import os
import re
import shutil
import subprocess

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

PROJECT_ROOT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "../")


def prepare(v: str):
    """Create a new release branch"""
    logger.info("Start to prepare release branch for version %s", v)
    _check_release_version(v)
    os.chdir(PROJECT_ROOT_DIR)
    branch = f"releases-{v}"
    try:
        subprocess.check_call(f"git checkout -b {branch}", shell=True)
        bump_version(version=v, l="all")
        subprocess.check_call("git add -u", shell=True)
        subprocess.check_call(f"git commit -m 'prepare release for {v}'", shell=True)
    except BaseException:
        logger.exception("Prepare branch failed")
        subprocess.check_call(f"git checkout - && git branch -D {branch}", shell=True)
        raise


def build(v: str):
    """version format: 0.5.1"""
    logger.info("Start to prepare release artifacts for version %s", v)
    _check_release_version(v)
    os.chdir(PROJECT_ROOT_DIR)
    if os.path.exists("dist"):
        shutil.rmtree("dist")
    os.mkdir("dist")
    branch = f"releases-{v}"
    # Check if branch exists, if not create it
    result = subprocess.run(
        f"git show-ref --verify --quiet refs/heads/{branch}",
        shell=True,
        capture_output=True,
    )
    if result.returncode == 0:
        # Branch exists, checkout
        subprocess.check_call(f"git checkout {branch}", shell=True)
    else:
        # Branch doesn't exist, create it
        subprocess.check_call(f"git checkout -b {branch}", shell=True)
    src_tar = f"apache-fory-{v}-src.tar.gz"
    _check_all_committed()
    _strip_unnecessary_license()
    subprocess.check_call(
        "git add LICENSE && git commit -m 'remove benchmark from license'", shell=True
    )
    subprocess.check_call(
        f"git archive --format=tar.gz "
        f"--output=dist/{src_tar} "
        f"--prefix=apache-fory-{v}-src/ {branch}",
        shell=True,
    )
    subprocess.check_call("git reset --hard HEAD~", shell=True)
    os.chdir("dist")
    logger.info("Start to generate signature")
    subprocess.check_call(
        f"gpg --armor --output {src_tar}.asc --detach-sig {src_tar}", shell=True
    )
    subprocess.check_call(f"sha512sum {src_tar} >{src_tar}.sha512", shell=True)
    verify(v)


def _check_release_version(v: str):
    assert v
    if "rc" in v:
        raise ValueError(
            "RC should only be contained in tag and svn directory, not in code"
        )


def _check_all_committed():
    proc = subprocess.run("git diff --quiet", capture_output=True, shell=True)
    result = proc.returncode
    if result != 0:
        raise Exception(
            f"There are some uncommitted files: {proc.stdout}, please commit it."
        )


def _strip_unnecessary_license():
    with open("LICENSE", "r") as f:
        lines = f.readlines()
    new_lines = []
    line_number = 0
    while line_number < len(lines):
        line = lines[line_number]
        if "fast-serialization" in line:
            line_number += 4
        elif "benchmark" in line:  # strip license in benchmark
            line_number += 1
        else:
            new_lines.append(line)
            line_number += 1
    text = "".join(new_lines)
    if lines != new_lines:
        with open("LICENSE", "w") as f:
            f.write(text)


def verify(v):
    src_tar = f"apache-fory-{v}-src.tar.gz"
    subprocess.check_call(f"gpg --verify {src_tar}.asc {src_tar}", shell=True)
    logger.info("Verified signature")
    subprocess.check_call(f"sha512sum --check {src_tar}.sha512", shell=True)
    logger.info("Verified checksum successfully")


def bump_version(**kwargs):
    new_version = kwargs["version"]
    langs = kwargs["l"]
    if langs == "all":
        langs = [
            "java",
            "python",
            "javascript",
            "scala",
            "rust",
            "kotlin",
            "cpp",
            "go",
            "dart",
            "csharp",
            "swift",
            "compiler",
        ]
    else:
        langs = langs.split(",")
    for lang in langs:
        if lang == "java":
            bump_java_version(_normalize_java_version(new_version))
        elif lang == "scala":
            _bump_version(
                "scala",
                "build.sbt",
                _normalize_java_version(new_version),
                _update_scala_version,
            )
        elif lang == "kotlin":
            _bump_version(
                "kotlin",
                "pom.xml",
                _normalize_java_version(new_version),
                _update_kotlin_version,
            )
        elif lang == "rust":
            bump_rust_version(new_version)
        elif lang == "python":
            bump_python_version(new_version)
        elif lang == "javascript":
            _bump_version(
                "javascript/packages/core",
                "package.json",
                _normalize_js_version(new_version),
                _update_js_version,
            )
            _bump_version(
                "javascript/packages/hps",
                "package.json",
                _normalize_js_version(new_version),
                _update_js_version,
            )
        elif lang == "cpp":
            bump_cpp_version(new_version)
        elif lang == "go":
            bump_go_version(new_version)
        elif lang == "dart":
            bump_dart_version(new_version)
        elif lang == "csharp":
            bump_csharp_version(new_version)
        elif lang == "swift":
            bump_swift_version(new_version)
        elif lang == "compiler":
            bump_compiler_version(new_version)
        else:
            raise NotImplementedError(f"Unsupported {lang}")


def _bump_version(path, file, new_version, func):
    os.chdir(os.path.join(PROJECT_ROOT_DIR, path))
    with open(file, "r") as f:
        lines = f.readlines()
    lines = func(lines, new_version) or lines
    text = "".join(lines)
    with open(file, "w") as f:
        f.write(text)


def bump_java_version(new_version):
    new_version = _normalize_java_version(new_version)
    for p in [
        "integration_tests/graalvm_tests",
        "integration_tests/jdk_compatibility_tests",
        "integration_tests/jpms_tests",
        "integration_tests/idl_tests/java",
        "benchmarks/java",
        "java/fory-core",
        "java/fory-format",
        "java/fory-simd",
        "java/fory-extensions",
        "java/fory-graalvm-feature",
        "java/fory-test-core",
        "java/fory-testsuite",
        "java/fory-latest-jdk-tests",
    ]:
        _bump_version(p, "pom.xml", new_version, _update_pom_parent_version)
    # mvn versions:set too slow
    # os.chdir(os.path.join(PROJECT_ROOT_DIR, "java"))
    # subprocess.check_output(
    #     f"mvn versions:set -DnewVersion={new_version}",
    #     shell=True,
    #     universal_newlines=True,
    # )
    _bump_version("java", "pom.xml", new_version, _update_parent_pom_version)


def bump_python_version(new_version):
    _bump_version("python/pyfory", "__init__.py", new_version, _update_python_version)
    _bump_version(
        "integration_tests/idl_tests/python",
        "pyproject.toml",
        new_version,
        _update_pyproject_version,
    )


def bump_rust_version(new_version):
    rust_version = _normalize_rust_version(new_version)
    _bump_version("rust", "Cargo.toml", rust_version, _update_rust_version)
    _bump_version(
        "benchmarks/rust",
        "Cargo.toml",
        rust_version,
        _update_cargo_package_version,
    )
    _bump_version(
        "integration_tests/idl_tests/rust",
        "Cargo.toml",
        rust_version,
        _update_cargo_package_version,
    )


def bump_cpp_version(new_version):
    for p in [
        "cpp",
        "benchmarks/cpp",
        "integration_tests/idl_tests/cpp",
    ]:
        _bump_version(p, "CMakeLists.txt", new_version, _update_cmake_project_version)


def bump_go_version(new_version):
    for p in [
        "benchmarks/go",
        "integration_tests/idl_tests/go",
    ]:
        _bump_version(p, "go.mod", new_version, _update_go_mod_version)


def bump_dart_version(new_version):
    for p in [
        "dart",
        "dart/packages/fory",
        "dart/packages/fory-test",
    ]:
        _bump_version(p, "pubspec.yaml", new_version, _update_pubspec_version)


def bump_compiler_version(new_version):
    _bump_version("compiler", "pyproject.toml", new_version, _update_pyproject_version)
    _bump_version(
        "compiler/fory_compiler",
        "__init__.py",
        new_version,
        _update_python_version,
    )


def bump_csharp_version(new_version):
    _bump_version(
        "csharp",
        "Directory.Build.props",
        new_version,
        _update_csharp_props_version,
    )
    _bump_version(
        "csharp",
        "README.md",
        new_version,
        _update_csharp_readme_package_version,
    )
    _bump_version(
        "docs/guide/csharp",
        "index.md",
        new_version,
        _update_csharp_readme_package_version,
    )


def bump_swift_version(new_version):
    _bump_version(
        "swift",
        "README.md",
        new_version,
        _update_swift_readme_dependency_version,
    )


def _update_pom_parent_version(lines, new_version):
    start_index, end_index = -1, -1
    for i, line in enumerate(lines):
        if "<parent>" in line:
            start_index = i
        if "</parent>" in line:
            end_index = i
            break
    assert start_index != -1
    assert end_index != -1
    for line_number in range(start_index, end_index):
        line = lines[line_number]
        if "version" in line:
            line = re.sub(
                r"(<version>)[^<>]+(</version>)", r"\g<1>" + new_version + r"\2", line
            )
            lines[line_number] = line


def _update_scala_version(lines, v):
    v = _normalize_java_version(v)
    for index, line in enumerate(lines):
        if "foryVersion = " in line:
            lines[index] = f'val foryVersion = "{v}"\n'
            break
    return lines


def _update_kotlin_version(lines, v):
    v = _normalize_java_version(v)
    return _update_pom_version(lines, v, "<artifactId>fory-kotlin</artifactId>")


def _update_parent_pom_version(lines, v):
    return _update_pom_version(lines, v, "<packaging>pom</packaging>")


def _update_pom_version(lines, v, prev):
    target_index = -1
    for index, line in enumerate(lines):
        if prev in line:
            target_index = index + 1
            break
    current_version_line = lines[target_index]
    # Find the start and end of the version number
    start = current_version_line.index("<version>") + len("<version>")
    end = current_version_line.index("</version>")
    # Replace the version number
    updated_version_line = current_version_line[:start] + v + current_version_line[end:]
    lines[target_index] = updated_version_line
    return lines


def _update_rust_version(lines, v):
    in_workspace_package = False
    in_workspace_dependencies = False
    for index, line in enumerate(lines):
        stripped = line.strip()
        if stripped == "[workspace.package]":
            in_workspace_package = True
            in_workspace_dependencies = False
            continue
        if stripped == "[workspace.dependencies]":
            in_workspace_dependencies = True
            in_workspace_package = False
            continue
        if stripped.startswith("[") and stripped.endswith("]"):
            in_workspace_package = False
            in_workspace_dependencies = False
        if in_workspace_package and stripped.startswith("version = "):
            lines[index] = f'version = "{v}"\n'
            continue
        if in_workspace_dependencies and re.match(r"\s*fory(-core|-derive)?\s*=", line):
            lines[index] = re.sub(
                r'(version\s*=\s*")([^"]+)(")',
                r"\g<1>" + v + r"\3",
                line,
            )
    return lines


def _update_python_version(lines, v: str):
    v = _normalize_python_version(v)
    for index, line in enumerate(lines):
        if "__version__ = " in line:
            lines[index] = f'__version__ = "{v}"\n'
            break


def _update_pyproject_version(lines, v: str):
    v = _normalize_python_version(v)
    in_project = False
    for index, line in enumerate(lines):
        stripped = line.strip()
        if stripped == "[project]":
            in_project = True
            continue
        if in_project and stripped.startswith("[") and stripped.endswith("]"):
            in_project = False
        if in_project and stripped.startswith("version ="):
            lines[index] = f'version = "{v}"\n'
            break
    return lines


def _update_cargo_package_version(lines, v: str):
    in_package = False
    for index, line in enumerate(lines):
        stripped = line.strip()
        if stripped == "[package]":
            in_package = True
            continue
        if in_package and stripped.startswith("[") and stripped.endswith("]"):
            in_package = False
        if in_package and stripped.startswith("version ="):
            lines[index] = f'version = "{v}"\n'
            break
    return lines


def _update_cmake_project_version(lines, v: str):
    cmake_version = _normalize_cmake_version(v)
    in_project = False
    for index, line in enumerate(lines):
        if re.search(r"^\s*project\(", line):
            in_project = True
        if in_project and "VERSION" in line:
            lines[index] = re.sub(
                r"(VERSION\s+)([0-9]+(?:\.[0-9]+){1,2})",
                r"\g<1>" + cmake_version,
                line,
            )
        if in_project and ")" in line:
            in_project = False
    return lines


def _update_go_mod_version(lines, v: str):
    go_version = _normalize_go_version(v)
    for index, line in enumerate(lines):
        if "github.com/apache/fory/go/fory" not in line:
            continue
        lines[index] = re.sub(
            r"(github.com/apache/fory/go/fory\s+)(v[^\s]+)",
            r"\g<1>" + go_version,
            line,
        )
    return lines


def _update_pubspec_version(lines, v: str):
    for index, line in enumerate(lines):
        if re.match(r"^version\s*:", line):
            lines[index] = f"version: {v}\n"
            continue
        if re.match(r"^\s*fory\s*:", line):
            prefix = re.match(r"^(\s*fory\s*:)\s*.*", line)
            if prefix:
                lines[index] = f"{prefix.group(1)} {v}\n"
    return lines


def _update_csharp_props_version(lines, v: str):
    for index, line in enumerate(lines):
        if "<Version>" not in line:
            continue
        lines[index] = re.sub(
            r"(<Version>)[^<]+(</Version>)",
            r"\g<1>" + v + r"\2",
            line,
        )
        return lines
    raise ValueError("No <Version> element found in csharp/Directory.Build.props")


def _update_csharp_readme_package_version(lines, v: str):
    for index, line in enumerate(lines):
        if "PackageReference" not in line or "Apache.Fory" not in line:
            continue
        lines[index] = re.sub(
            r'(<PackageReference\s+Include="Apache\.Fory"\s+Version=")[^"]+(")',
            r"\g<1>" + v + r"\2",
            line,
        )
        return lines
    raise ValueError("No Apache.Fory PackageReference version snippet found")


def _update_swift_readme_dependency_version(lines, v: str):
    for index, line in enumerate(lines):
        if "https://github.com/apache/fory.git" not in line:
            continue
        lines[index] = re.sub(
            r'(\.package\(url:\s*"https://github\.com/apache/fory\.git",\s*from:\s*")[^"]+("\))',
            r"\g<1>" + v + r"\2",
            line,
        )
        return lines
    raise ValueError("No Swift Package dependency snippet for apache/fory.git found")


def _normalize_python_version(v: str) -> str:
    v = v.strip()
    v = re.sub(r"(?i)-?snapshot$", ".dev0", v)
    v = re.sub(r"(?i)-dev(\d+)$", r".dev\1", v)
    v = re.sub(r"(?i)-dev$", ".dev0", v)
    v = v.replace("-alpha", "a")
    v = v.replace("-beta", "b")
    v = v.replace("-rc", "rc")
    v = v.replace("-", "")
    return v


def _normalize_java_version(v: str) -> str:
    v = v.strip()
    if re.search(r"(?i)-snapshot$", v):
        return re.sub(r"(?i)-snapshot$", "-SNAPSHOT", v)
    if re.search(r"(?i)(\.dev\d*|-dev\d*)$", v):
        base = re.sub(r"(?i)(\.dev\d*|-dev\d*)$", "", v)
        return f"{base}-SNAPSHOT"
    return v


def _normalize_go_version(v: str) -> str:
    v = v.strip()
    if v.startswith("v"):
        v = v[1:]
    v = re.sub(r"-(alpha|beta|rc)(\d+)$", r"-\1.\2", v)
    if re.search(r"(?i)-(alpha|beta)\.0$", v):
        return f"v{v}"
    if re.search(r"(?i)-(alpha|beta|rc)\.\d+$", v):
        return f"v{v}"
    if re.search(r"(?i)-pre$", v):
        return f"v{v}"
    dev_match = re.search(r"(?i)(?:-dev|\.dev)(\d+)$", v)
    if dev_match:
        base = re.sub(r"(?i)(?:-dev|\.dev)\d+$", "", v)
        return f"v{base}-alpha.{dev_match.group(1)}"
    if re.search(r"(?i)(-snapshot|\.dev|-dev)$", v):
        base = re.sub(r"(?i)(-snapshot|\.dev|-dev)$", "", v)
        return f"v{base}-alpha.0"
    return f"v{v}"


def _normalize_cmake_version(v: str) -> str:
    v = v.strip()
    if v.startswith("v"):
        v = v[1:]
    v = re.split(r"[-+]", v, maxsplit=1)[0]
    return v


def _update_js_version(lines, v: str):
    v = _normalize_js_version(v)
    for index, line in enumerate(lines):
        if "version" in line:
            # "version": "0.5.9-beta"
            for x in ["-alpha", "-beta", "-rc"]:
                if x in v and v.split(x)[-1].isdigit():
                    v = v.replace(x, x + ".")
            lines[index] = f'  "version": "{v}",\n'
            break


def _normalize_js_version(v: str) -> str:
    v = v.strip()
    v = re.sub(r"-(alpha|beta|rc)(\d+)$", r"-\1.\2", v)
    dev_match = re.search(r"(?i)(?:-dev|\\.dev)(\\d+)$", v)
    if dev_match:
        v = re.sub(r"(?i)(?:-dev|\\.dev)\\d+$", f"-alpha.{dev_match.group(1)}", v)
        return v
    if re.search(r"(?i)(-snapshot|\\.dev|-dev)$", v):
        v = re.sub(r"(?i)(-snapshot|\\.dev|-dev)$", "-alpha.0", v)
    return v


def _normalize_rust_version(v: str) -> str:
    v = v.strip()
    v = re.sub(r"-(alpha|beta|rc)(\d+)$", r"-\1.\2", v)
    if re.search(r"(?i)-(alpha|beta)\.0$", v):
        return v
    if re.search(r"(?i)-(alpha|beta|rc)\.\d+$", v):
        return v
    if re.search(r"(?i)-pre$", v):
        return v
    dev_match = re.search(r"(?i)(?:-dev|\\.dev)(\\d+)$", v)
    if dev_match:
        base = re.sub(r"(?i)(?:-dev|\\.dev)\\d+$", "", v)
        return f"{base}-alpha.{dev_match.group(1)}"
    if re.search(r"(?i)(-snapshot|\\.dev|-dev)$", v):
        base = re.sub(r"(?i)(-snapshot|\\.dev|-dev)$", "", v)
        return f"{base}-alpha.0"
    return v


def _parse_args():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.set_defaults(func=parser.print_help)
    subparsers = parser.add_subparsers()
    bump_version_parser = subparsers.add_parser(
        "bump_version",
        description="Bump version",
    )
    bump_version_parser.add_argument("-version", type=str, help="new version")
    bump_version_parser.add_argument("-l", type=str, help="language")
    bump_version_parser.set_defaults(func=bump_version)

    prepare_parser = subparsers.add_parser(
        "prepare",
        description="Prepare release branch",
    )
    prepare_parser.add_argument("-v", type=str, help="new version")
    prepare_parser.set_defaults(func=prepare)

    release_parser = subparsers.add_parser(
        "build",
        description="Build release artifacts",
    )
    release_parser.add_argument("-v", type=str, help="new version")
    release_parser.set_defaults(func=build)

    verify_parser = subparsers.add_parser(
        "verify",
        description="Verify release artifacts",
    )
    verify_parser.add_argument("-v", type=str, help="new version")
    verify_parser.set_defaults(func=verify)

    args = parser.parse_args()
    arg_dict = dict(vars(args))
    del arg_dict["func"]
    print(arg_dict)
    args.func(**arg_dict)


if __name__ == "__main__":
    _parse_args()
