#!/usr/bin/env python3
# Copyright (c) 2025 Huawei Device Co., Ltd.
# Licensed 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 shutil
import subprocess
import tarfile
from datetime import datetime
from pathlib import Path
import struct
import argparse
import zlib
import sys

sys.path.insert(0, str(Path(__file__).parent / "../compiler"))
from taihe.utils.resources import (
    BUNDLE_PYTHON_RUNTIME_DIR_NAME,
    DeploymentMode,
    ResourceLocator,
    ResourceType,
)


SYSTEM_TYPES = ["linux-x86_64", "windows-x86_64", "darwin-arm64", "darwin-x86_64"]


class TaiheBuilder:
    def __init__(self, project_dir: Path, system: str, version: str):
        self.system = system
        self.version = version

        self.project_dir = project_dir
        self.compiler_dir = self.project_dir / "compiler"
        self.resources = ResourceLocator.detect()

        self.dist = self.project_dir / "dist" / f"taihe-{system}" / "taihe"
        if self.dist.exists():
            shutil.rmtree(self.dist)
        self.dist_pyrt_dir = self.dist / "lib" / BUNDLE_PYTHON_RUNTIME_DIR_NAME
        self.dist_bin_dir = self.dist / "bin"
        self.dist_resources = ResourceLocator(DeploymentMode.BUNDLE, self.dist)

    def copy_resources(self):
        for ty in ResourceType:
            if ty.is_packagable():
                shutil.copytree(
                    self.resources.get(ty),
                    self.dist_resources.get(ty),
                    dirs_exist_ok=True,
                    ignore=shutil.ignore_patterns("build", "generated"),
                )

    def install_compiler(self):
        print(f"Copying the compiler for {self.system}")
        site_packages_dir = next(self.dist_pyrt_dir.glob("lib/python*/site-packages"))

        # TODO rm dummy scripts / fix shebang paths
        uv_args = [
            "uv",
            "pip",
            "install",
            f"--target={site_packages_dir}",
            "--",
            str(self.compiler_dir),
        ]
        subprocess.run(uv_args, check=True, cwd=self.compiler_dir)

    def extract_python(self):
        python_tar = (
            self.resources.get(ResourceType.DEV_PYTHON_BUILD)
            / f"{self.system}-python.tar.gz"
        )
        # First, extract to "dist/lib".
        # It creates a sole directory named "python"
        parent_dir = self.dist_pyrt_dir.parent
        with tarfile.open(python_tar, "r:gz") as tar:
            tar.extractall(parent_dir, filter="tar")
        # Next, rename "dist/lib/python" to "dist/lib/pyrt"
        old_dir = parent_dir / "python"
        if self.dist_pyrt_dir.exists():
            self.dist_pyrt_dir.rmdir()
        old_dir.rename(self.dist_pyrt_dir)

    def create_script(self, binary: str, pymod: str):
        # TODO use absolute script path
        is_windows = self.system.startswith("windows")
        prog = self.dist_bin_dir / (f"{binary}.bat" if is_windows else binary)
        if is_windows:
            content = (
                "@echo off\n"
                "set TAIHE_ROOT=%~dp0..\n"
                f'%TAIHE_ROOT%\\lib\\pyrt\\bin\\python3.11.exe -m "{pymod}" %*\n'
            )
        else:
            content = (
                "#!/bin/bash -eu\n"
                'export TAIHE_ROOT="$(realpath $(dirname "$0")/..)"\n'
                f'exec "$TAIHE_ROOT/lib/pyrt/bin/python3" -m \'{pymod}\' "$@"\n'
            )
        prog.write_text(content)
        prog.chmod(0o755)

    def create_taihec_script(self):
        self.dist_bin_dir.mkdir()
        print(f"Creating taihec script for {self.system}")
        if self.system.startswith("linux"):
            self.create_script("taihe-tryit", "taihe.cli.tryit")
        self.create_script("taihec", "taihe.cli.compiler")

    def write_version(self):
        print(f"Writing version information for {self.system}")
        version_path = self.dist / "version.txt"

        git_commit = subprocess.run(
            ["git", "rev-parse", "HEAD"],
            capture_output=True,
            text=True,
            check=True,
            cwd=self.project_dir,
        ).stdout.strip()

        git_message = (
            subprocess.run(
                ["git", "log", "-1", "--pretty=%B"],
                capture_output=True,
                text=True,
                check=True,
                cwd=self.project_dir,
            )
            .stdout.splitlines()[0]
            .strip()
        )

        with open(version_path, "w") as version_file:
            version_file.write(f"System        : {self.system}\n")
            version_file.write(f"Version       : {self.version}\n")
            version_file.write(f"Last Commit   : {git_commit}\n")
            version_file.write(f"Commit Message: {git_message}\n")

    def create_package(self):
        print(f"Creating package for {self.system}")
        package_name = (
            f"taihe-{self.system}-{self.version}-{datetime.now().strftime('%Y%m%d')}"
        )
        tar_path = self.project_dir / f"{package_name}.tar"
        zip_path = self.project_dir / f"{package_name}.tar.gz"

        def reset_tarinfo(tarinfo: tarfile.TarInfo):
            tarinfo.uid = 0
            tarinfo.gid = 0
            tarinfo.uname = ""
            tarinfo.gname = ""
            tarinfo.mtime = 1704067200
            return tarinfo

        try:
            with tarfile.open(tar_path, "w:", format=tarfile.GNU_FORMAT) as tar:
                for file_path in sorted(self.dist.rglob("*")):
                    if file_path.is_file():
                        arcname = f"{self.dist.name}/{file_path.relative_to(self.dist)}"
                        tar.add(file_path, arcname=arcname, filter=reset_tarinfo)
            with open(tar_path, "rb") as tar_file, open(zip_path, "wb") as zip_file:
                zip_file.write(b"\x1f\x8b\x08\x00\x00\x00\x00\x00\x00\x03")
                compressor = zlib.compressobj(
                    9,
                    zlib.DEFLATED,
                    -zlib.MAX_WBITS,
                    zlib.DEF_MEM_LEVEL,
                    0,
                )
                data = tar_file.read()
                compressed = compressor.compress(data) + compressor.flush()
                zip_file.write(compressed)
                crc = zlib.crc32(data)
                size = len(data)
                zip_file.write(struct.pack("<II", crc & 0xFFFFFFFF, size & 0xFFFFFFFF))
        finally:
            if tar_path.exists():
                tar_path.unlink()

        print(f"Created package for {self.system}: {zip_path}")

    def build(self):
        self.extract_python()
        self.install_compiler()
        self.copy_resources()
        self.create_taihec_script()
        self.write_version()
        self.create_package()


def main():
    parser = argparse.ArgumentParser(description="Build the Taihe project.")
    parser.add_argument(
        "--system",
        choices=SYSTEM_TYPES,
        nargs="*",
        default=None,
        help="Specify the system type to build for. Can be used multiple times.",
    )
    args = parser.parse_args()

    project_dir = Path(__file__).parent.parent

    # TODO detect version use uv / hatch
    try:
        version = subprocess.run(
            ["git", "describe", "--tags", "--abbrev=0"],
            capture_output=True,
            text=True,
            check=True,
            cwd=project_dir,
        ).stdout.strip()
    except subprocess.CalledProcessError:
        short_hash = subprocess.run(
            ["git", "rev-parse", "--short=8", "HEAD"],
            capture_output=True,
            text=True,
            check=True,
            cwd=project_dir,
        ).stdout.strip()
        version = f"TRUNK-{short_hash}"

    if args.system is None:
        systems = SYSTEM_TYPES
    else:
        systems = args.system

    print(f"Version is {version}")
    for system in systems:
        print(f"Building for {system}")
        builder = TaiheBuilder(project_dir, system, version)
        builder.build()


if __name__ == "__main__":
    main()