import requests
import os
import ast
import re
import argparse
import sys
import subprocess
from importlib import import_module
import platform
from pathlib import Path
import sys
import zipfile
import shutil
import stat


here = Path(__file__).parent
executable = sys.executable

# - detect mode
mode = None
if os.environ.get("COLAB_GPU"):
    mode = "colab"
elif "python_embeded" in executable:
    mode = "embeded"
elif ".venv" in executable:
    mode = "venv"


if mode == None:
    mode = "unknown"

# region ansi
# ANSI escape sequences for text styling
ANSI_FORMATS = {
    "reset": "\033[0m",
    "bold": "\033[1m",
    "dim": "\033[2m",
    "italic": "\033[3m",
    "underline": "\033[4m",
    "blink": "\033[5m",
    "reverse": "\033[7m",
    "strike": "\033[9m",
}

ANSI_COLORS = {
    "black": "\033[30m",
    "red": "\033[31m",
    "green": "\033[32m",
    "yellow": "\033[33m",
    "blue": "\033[34m",
    "magenta": "\033[35m",
    "cyan": "\033[36m",
    "white": "\033[37m",
    "bright_black": "\033[30;1m",
    "bright_red": "\033[31;1m",
    "bright_green": "\033[32;1m",
    "bright_yellow": "\033[33;1m",
    "bright_blue": "\033[34;1m",
    "bright_magenta": "\033[35;1m",
    "bright_cyan": "\033[36;1m",
    "bright_white": "\033[37;1m",
    "bg_black": "\033[40m",
    "bg_red": "\033[41m",
    "bg_green": "\033[42m",
    "bg_yellow": "\033[43m",
    "bg_blue": "\033[44m",
    "bg_magenta": "\033[45m",
    "bg_cyan": "\033[46m",
    "bg_white": "\033[47m",
    "bg_bright_black": "\033[40;1m",
    "bg_bright_red": "\033[41;1m",
    "bg_bright_green": "\033[42;1m",
    "bg_bright_yellow": "\033[43;1m",
    "bg_bright_blue": "\033[44;1m",
    "bg_bright_magenta": "\033[45;1m",
    "bg_bright_cyan": "\033[46;1m",
    "bg_bright_white": "\033[47;1m",
}


def apply_format(text, *formats):
    """Apply ANSI escape sequences for the specified formats to the given text."""
    formatted_text = text
    for format in formats:
        formatted_text = f"{ANSI_FORMATS.get(format, '')}{formatted_text}{ANSI_FORMATS.get('reset', '')}"
    return formatted_text


def apply_color(text, color=None, background=None):
    """Apply ANSI escape sequences for the specified color and background to the given text."""
    formatted_text = text
    if color:
        formatted_text = f"{ANSI_COLORS.get(color, '')}{formatted_text}{ANSI_FORMATS.get('reset', '')}"
    if background:
        formatted_text = f"{ANSI_COLORS.get(background, '')}{formatted_text}{ANSI_FORMATS.get('reset', '')}"
    return formatted_text


def print_formatted(text, *formats, color=None, background=None, **kwargs):
    """Print the given text with the specified formats, color, and background."""
    formatted_text = apply_format(text, *formats)
    formatted_text = apply_color(formatted_text, color, background)
    file = kwargs.get("file", sys.stdout)
    print(
        apply_color(apply_format("[mtb install] ", "bold"), color="yellow"),
        formatted_text,
        file=file,
    )


# endregion

try:
    import requirements
except ImportError:
    print_formatted("Installing requirements-parser...", "italic", color="yellow")
    subprocess.check_call(
        [sys.executable, "-m", "pip", "install", "requirements-parser"]
    )
    import requirements

    print_formatted("Done.", "italic", color="green")

try:
    from tqdm import tqdm
except ImportError:
    print_formatted("Installing tqdm...", "italic", color="yellow")
    subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "tqdm"])
    from tqdm import tqdm
import importlib


pip_map = {
    "onnxruntime-gpu": "onnxruntime",
    "opencv-contrib": "cv2",
    "tb-nightly": "tensorboard",
    "protobuf": "google.protobuf",
    # Add more mappings as needed
}


def is_pipe():
    try:
        mode = os.fstat(0).st_mode
        return (
            stat.S_ISFIFO(mode)
            or stat.S_ISREG(mode)
            or stat.S_ISBLK(mode)
            or stat.S_ISSOCK(mode)
        )
    except OSError:
        return False


# Get the version from __init__.py
def get_local_version():
    init_file = os.path.join(os.path.dirname(__file__), "__init__.py")
    if os.path.isfile(init_file):
        with open(init_file, "r") as f:
            tree = ast.parse(f.read())
            for node in ast.walk(tree):
                if isinstance(node, ast.Assign):
                    for target in node.targets:
                        if (
                            isinstance(target, ast.Name)
                            and target.id == "__version__"
                            and isinstance(node.value, ast.Str)
                        ):
                            return node.value.s
    return None


def download_file(url, file_name):
    with requests.get(url, stream=True) as response:
        response.raise_for_status()
        total_size = int(response.headers.get("content-length", 0))
        with open(file_name, "wb") as file, tqdm(
            desc=file_name.stem,
            total=total_size,
            unit="B",
            unit_scale=True,
            unit_divisor=1024,
        ) as progress_bar:
            for chunk in response.iter_content(chunk_size=8192):
                file.write(chunk)
                progress_bar.update(len(chunk))


def get_requirements(path: Path):
    with open(path.resolve(), "r") as requirements_file:
        requirements_txt = requirements_file.read()

    try:
        parsed_requirements = requirements.parse(requirements_txt)
    except AttributeError:
        print_formatted(
            f"Failed to parse {path}. Please make sure the file is correctly formatted.",
            "bold",
            color="red",
        )

        return

    return parsed_requirements


def try_import(requirement):
    dependency = requirement.name.strip()
    import_name = pip_map.get(dependency, dependency)
    installed = False

    pip_name = dependency
    if specs := requirement.specs:
        pip_name += "".join(specs[0])

    try:
        import_module(import_name)
        print_formatted(
            f"Package {pip_name} already installed (import name: '{import_name}').",
            "bold",
            color="green",
        )
        installed = True
    except ImportError:
        pass

    return (installed, pip_name, import_name)


def import_or_install(requirement, dry=False):
    installed, pip_name, import_name = try_import(requirement)

    if not installed:
        print_formatted(f"Installing package {pip_name}...", "italic", color="yellow")
        if dry:
            print_formatted(
                f"Dry-run: Package {pip_name} would be installed (import name: '{import_name}').",
                color="yellow",
            )
        else:
            try:
                subprocess.check_call(
                    [sys.executable, "-m", "pip", "install", pip_name]
                )
                print_formatted(
                    f"Package {pip_name} installed successfully using pip package name  (import name: '{import_name}')",
                    "bold",
                    color="green",
                )
            except subprocess.CalledProcessError as e:
                print_formatted(
                    f"Failed to install package {pip_name} using pip package name  (import name: '{import_name}'). Error: {str(e)}",
                    "bold",
                    color="red",
                )


# Install dependencies from requirements.txt
def install_dependencies(dry=False):
    parsed_requirements = get_requirements(here / "requirements.txt")
    if not parsed_requirements:
        return
    print_formatted(
        "Installing dependencies from requirements.txt...", "italic", color="yellow"
    )

    for requirement in parsed_requirements:
        import_or_install(requirement, dry=dry)

    if mode == "venv":
        parsed_requirements = get_requirements(here / "requirements-wheels.txt")
        if not parsed_requirements:
            return
        for requirement in parsed_requirements:
            import_or_install(requirement, dry=dry)


if __name__ == "__main__":
    full = False
    if is_pipe():
        print_formatted("Pipe detected, full install...", color="green")
        # we clone our repo
        url = "https://github.com/melmass/comfy_mtb.git"
        clone_dir = here / "custom_nodes" / "comfy_mtb"
        if not clone_dir.exists():
            clone_dir.parent.mkdir(parents=True, exist_ok=True)
            print_formatted(f"Cloning {url} to {clone_dir}", "italic", color="yellow")
            subprocess.check_call(["git", "clone", "--recursive", url, clone_dir])

        # os.chdir(clone_dir)
        here = clone_dir
        full = True

    if len(sys.argv) == 1:
        print_formatted(
            "No arguments provided, doing a full install/update...",
            "italic",
            color="yellow",
        )

        full = True

    # Parse command-line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--wheels", "-w", action="store_true", help="Install wheel dependencies"
    )
    parser.add_argument(
        "--requirements", "-r", action="store_true", help="Install requirements.txt"
    )
    parser.add_argument(
        "--dry",
        action="store_true",
        help="Print what will happen without doing it (still making requests to the GH Api)",
    )

    # parser.add_argument(
    #     "--version",
    #     default=get_local_version(),
    #     help="Version to check against the GitHub API",
    # )

    args = parser.parse_args()

    wheels_directory = here / "wheels"
    print_formatted(f"Detected environment: {apply_color(mode,'cyan')}")

    # Install dependencies from requirements.txt
    # if args.requirements or mode == "venv":
    install_dependencies(dry=args.dry)

    if (not args.wheels and mode not in ["colab", "embeded"]) and not full:
        print_formatted(
            "Skipping wheel installation. Use --wheels to install wheel dependencies. (only needed for Comfy embed)",
            "italic",
            color="yellow",
        )
        sys.exit()

    if mode in ["colab", "embeded"]:
        print_formatted(
            f"Downloading and installing release wheels since we are in a Comfy {apply_color(mode,'cyan')} environment",
        )
    if full:
        print_formatted(
            f"Downloading and installing release wheels since no arguments where provided"
        )

    # - Check the env before proceeding.
    missing_wheels = False
    parsed_requirements = get_requirements(here / "requirements-wheels.txt")
    if parsed_requirements:
        for requirement in parsed_requirements:
            installed, pip_name, import_name = try_import(requirement)
            if not installed:
                missing_wheels = True
                break

    if not missing_wheels:
        print_formatted(
            f"All required wheels are already installed.", "italic", color="green"
        )
        sys.exit()

    # Fetch the JSON data from the GitHub API URL
    owner = "melmass"
    repo = "comfy_mtb"
    # version = args.version
    current_platform = platform.system().lower()

    # Get the tag version from the GitHub API
    tag_url = f"https://api.github.com/repos/{owner}/{repo}/releases/latest"
    response = requests.get(tag_url)
    if response.status_code == 404:
        # print_formatted(
        #     f"Tag version '{apply_color(version,'cyan')}' not found for {owner}/{repo} repository."
        # )
        print_formatted("Error retrieving the release assets.", color="red")
        sys.exit()

    tag_data = response.json()
    tag_name = tag_data["name"]

    # # Compare the local and tag versions
    # if version and tag_name:
    #     if re.match(r"v?(\d+(\.\d+)+)", version) and re.match(
    #         r"v?(\d+(\.\d+)+)", tag_name
    #     ):
    #         version_parts = [int(part) for part in version.lstrip("v").split(".")]
    #         tag_version_parts = [int(part) for part in tag_name.lstrip("v").split(".")]

    #         if version_parts > tag_version_parts:
    #             print_formatted(
    #                 f"Local version ({version}) is greater than the release version ({tag_name}).",
    #                 "bold",
    #                 "yellow",
    #             )
    #             sys.exit()

    # Download the assets for the given version
    matching_assets = [
        asset for asset in tag_data["assets"] if current_platform in asset["name"]
    ]
    if not matching_assets:
        print_formatted(
            f"Unsupported operating system: {current_platform}", color="yellow"
        )

    wheels_directory.mkdir(exist_ok=True)

    for asset in matching_assets:
        asset_name = asset["name"]
        asset_download_url = asset["browser_download_url"]
        print_formatted(f"Downloading asset: {asset_name}", color="yellow")
        asset_dest = wheels_directory / asset_name
        download_file(asset_download_url, asset_dest)

        # - Unzip to wheels dir
        whl_files = []
        with zipfile.ZipFile(asset_dest, "r") as zip_ref:
            for item in tqdm(zip_ref.namelist(), desc="Extracting", unit="file"):
                if item.endswith(".whl"):
                    item_basename = os.path.basename(item)
                    target_path = wheels_directory / item_basename
                    with zip_ref.open(item) as source, open(
                        target_path, "wb"
                    ) as target:
                        whl_files.append(target_path)
                        shutil.copyfileobj(source, target)

        print_formatted(
            f"Wheels extracted for {current_platform} to the '{wheels_directory}' directory.",
            "bold",
            color="green",
        )

        if whl_files:
            for whl_file in tqdm(whl_files, desc="Installing", unit="package"):
                whl_path = wheels_directory / whl_file

                # check if installed
                try:
                    whl_dep = whl_path.name.split("-")[0]
                    import_name = pip_map.get(whl_dep, whl_dep)
                    import_module(import_name)
                    tqdm.write(
                        f"Package {import_name} already installed, skipping wheel installation.",
                    )
                    continue
                except ImportError:
                    if args.dry:
                        tqdm.write(
                            f"Dry-run: Package {whl_path.name} would be installed.",
                        )
                        continue

                    tqdm.write("Installing wheel: " + whl_path.name)

                    subprocess.check_call(
                        [
                            sys.executable,
                            "-m",
                            "pip",
                            "install",
                            whl_path.resolve().as_posix(),
                        ]
                    )

            print_formatted("Wheels installation completed.", color="green")
        else:
            print_formatted("No .whl files found. Nothing to install.", color="yellow")