# Copyright (c) 2025, UChicago Argonne, LLC
# BSD OPEN SOURCE LICENSE. Full license can be found in LICENSE
# Copyright (c) 2025, UChicago Argonne, LLC
# BSD OPEN SOURCE LICENSE. Full license can be found in LICENSE
# This portion of the script initializes the plugin, making it known to QGIS.
import numpy as np
import os
import shutil
import subprocess
import sys
from importlib.util import find_spec
from os.path import join, isdir
from pathlib import Path
from qgis.core import QgsApplication
from qgis.core import QgsMessageLog
from time import sleep
from typing import List


class download_all:
    must_remove = [
        "numpy",
        "scipy",
        "pandas",
        "charset_normalizer",
        "click_plugins",
        "click",
        "certifi",
        "cligj",
        "colorama",
        "fiona",
        "pyproj",
        "pytz",
        "requests",
        "rtree",
        "setuptools",
        "shapely",
        "six",
        "tzdata",
        "zipp",
        "attr",
        "attrs",
        "dateutil",
        "python_dateutil",
        "idna",
        "importlib_metadata",
        "pyaml",
        "urllib3",
        "packaging",
        "cpuinfo",
    ]

    def __init__(self, filename):
        pth = os.path.dirname(__file__)
        self.file = join(pth, filename)

        self.pth = str(Path(__file__).parent / "packages")
        self._file_done = join(self.pth, filename)
        self.no_ssl = False

    def install(self):
        # self.adapt_aeq_version()
        if Path(self._file_done).exists():
            return

        with open(self.file, "r") as fl:
            lines = fl.readlines()

        reps = []
        self.install_package("uv")
        failed_packages = []
        for line in lines:
            msg, failure = self.install_package(line.strip())
            reps.extend(msg)
            if failure:
                failed_packages.append(line.strip())

        for _ in range(100):
            if not failed_packages:
                break
            print("Retrying failed installation for:", failed_packages)
            for package in failed_packages:
                msg, failure = self.install_package(package)
                reps.extend(msg)
                if not failure:
                    failed_packages.remove(package)
            sleep(1)

        if failed_packages:
            print("The following packages failed to install:", failed_packages)
            QgsMessageLog.logMessage(
                "The following packages failed to install after 100 tries: " + ", ".join(failed_packages),
            )

        with open(self._file_done, "w") as _:
            pass
        return reps

    def install_package(self, package) -> (List[str], bool):
        arg = ""
        if any(x in package.lower() for x in ["polaris", "aequilibrae", "openmatrix"]):
            arg = " --no-deps"

        spec = find_spec("uv")
        is_uv = "" if spec is None else "uv"

        install_command = f'-m {is_uv} pip install {package} {arg} --target "{self.pth}"'

        command = f'"{self.find_python()}" {install_command}'
        print(command)

        if not self.no_ssl:
            reps = self.execute(command)

        if self.no_ssl or (
            "because the ssl module is not available" in "".join(reps).lower() and sys.platform == "win32"
        ):
            command = f"python {install_command}"
            print(command)
            reps = self.execute(command)
            self.no_ssl = True

        for line in reps:
            QgsMessageLog.logMessage(str(line))

        msgs = "".join([str(line).lower() for line in reps])
        failures = ["denied", "error", "failed"]

        return reps, any(failure in msgs for failure in failures)

    def execute(self, command):
        lines = []
        lines.append(command)
        with subprocess.Popen(
            command,
            shell=True,
            stdout=subprocess.PIPE,
            stdin=subprocess.DEVNULL,
            stderr=subprocess.STDOUT,
            universal_newlines=True,
        ) as proc:
            lines.extend(proc.stdout.readlines())
            exit_code = proc.wait()
            if exit_code != 0:
                print("Error code: ", exit_code)
        return lines

    def find_python(self):
        sys_exe = Path(sys.executable)
        if sys.platform == "linux" or sys.platform == "linux2":
            # Unlike other platforms, linux uses the system python, lets see if we can guess it
            if Path("/usr/bin/python3").exists():
                return "/usr/bin/python3"
            if Path("/usr/bin/python").exists():
                return "/usr/bin/python"
            # If that didn't work, it also has a valid sys.executable (unlike other platforms)
            python_exe = sys_exe

        # On mac/windows sys.executable returns '/Applications/QGIS.app/Contents/MacOS/QGIS' or
        # 'C:\\Program Files\\QGIS 3.30.0\\bin\\qgis-bin.exe' respectively so we need to explore in that area
        # of the filesystem
        elif sys.platform == "darwin":
            python_exe = sys_exe.parent / "bin" / "python3"
        elif sys.platform == "win32":
            python_exe = Path(sys.base_prefix) / "python3.exe"

        if not python_exe.exists():
            raise FileExistsError("Can't find a python executable to use")
        print(python_exe)
        return python_exe

    def adapt_aeq_version(self):
        if int(np.__version__.split(".")[1]) >= 22:
            Path(self.file).unlink(missing_ok=True)
            shutil.copyfile(self._file, self.file)
            return

        with open(self._file, "r") as fl:
            cts = [c.rstrip() for c in fl.readlines()]

        with open(self.file, "w") as fl:
            for c in cts:
                if "aequilibrae" in c:
                    c = c + ".dev0"
                fl.write(f"{c}\n")

    def clean_packages(self):
        QgsMessageLog.logMessage("Starting package cleaning")
        for fldr in os.listdir(self.pth):
            for pkg in self.must_remove:
                if pkg.lower() in fldr.lower():
                    if isdir(join(self.pth, fldr)):
                        QgsMessageLog.logMessage(f"Cleaning {pkg}")
                        shutil.rmtree(join(self.pth, fldr))
                        QgsMessageLog.logMessage(f"Done cleaning {pkg}")


def install_requirements():
    is_ci = "CI_POLARIS_STUDIO_ARTIFACT_DIR" in os.environ
    if is_ci:
        # If running in CI, we don't install packages
        return
    download_all("install_requirements.txt").install()
    download_all("polaris_studio_requirements.txt").install()
    download_all("").clean_packages()
    QgsMessageLog.logMessage("Finished installing QPolaris")
