import glob
import json
import os
import platform
import subprocess
import sys
import tempfile
from types import ModuleType
from typing import TYPE_CHECKING, Optional, Union
import uuid

import psutil

if TYPE_CHECKING:
    try:
        import ensight
    except ImportError:
        from ansys.api.pyensight import ensight_api


class Omniverse:
    """Provides the ``ensight.utils.omniverse`` interface.

    The omniverse class methods provide an interface between an EnSight session
    and an Omniverse instance.  See :ref:`omniverse_info` for additional details.

    Parameters
    ----------
    interface: Union["ensight_api.ensight", "ensight"]
        Entity that provides the ``ensight`` namespace. In the case of
        EnSight Python, the ``ensight`` module is passed. In the case
        of PyEnSight, ``Session.ensight`` is passed.

    Notes
    -----
    This interface is only available when using pyensight (they do not work with
    the ensight Python interpreter) and the module must be used in an interpreter
    that includes the Omniverse Python modules (e.g. omni and pxr).  Only a single
    Omniverse connection can be established within a single pyensight session.

    Examples
    --------

    >>> from ansys.pyensight.core import LocalLauncher
    >>> session = LocalLauncher().start()
    >>> ov = session.ensight.utils.omniverse
    >>> ov.create_connection(r"D:\Omniverse\Example")
    >>> ov.update()
    >>> ov.close_connection()

    """

    def __init__(self, interface: Union["ensight_api.ensight", "ensight"]):
        self._ensight = interface
        self._server_pid: Optional[int] = None
        self._interpreter: str = ""
        self._status_filename: str = ""

    @staticmethod
    def find_kit_filename(fallback_directory: Optional[str] = None) -> Optional[str]:
        """
        Use a combination of the current omniverse application and the information
        in the local .nvidia-omniverse/config/omniverse.toml file to come up with
        the pathname of a kit executable suitable for hosting another copy of the
        ansys.geometry.server kit.

        Returns
        -------
        Optional[str]
            The pathname of a kit executable or None

        """
        # parse the toml config file for the location of the installed apps
        try:
            import tomllib
        except ModuleNotFoundError:
            import pip._vendor.tomli as tomllib

        homedir = os.path.expanduser("~")
        ov_config = os.path.join(homedir, ".nvidia-omniverse", "config", "omniverse.toml")
        if not os.path.exists(ov_config):
            return None
        # read the Omniverse configuration toml file
        with open(ov_config, "r") as ov_file:
            ov_data = ov_file.read()
        config = tomllib.loads(ov_data)
        appdir = config.get("paths", {}).get("library_root", fallback_directory)

        # If we are running inside an Omniverse app, use that information
        try:
            import omni.kit.app

            # get the current application
            app = omni.kit.app.get_app()
            app_name = app.get_app_filename().split(".")[-1]
            app_version = app.get_app_version().split("-")[0]
            # and where it is installed
            appdir = os.path.join(appdir, f"{app_name}-{app_version}")
        except ModuleNotFoundError:
            # Names should be like: "C:\\Users\\foo\\AppData\\Local\\ov\\pkg\\create-2023.2.3\\launcher.toml"
            target = None
            target_version = None
            for d in glob.glob(os.path.join(appdir, "*", "launcher.toml")):
                test_dir = os.path.dirname(d)
                # the name will be something like "create-2023.2.3"
                name = os.path.basename(test_dir).split("-")
                if len(name) != 2:
                    continue
                if name[0] not in ("kit", "create", "view"):
                    continue
                if (target_version is None) or (name[1] > target_version):
                    target = test_dir
                    target_version = name[1]
            if target is None:
                return None
            appdir = target

        # Windows: 'kit.bat' in '.' or 'kit' followed by 'kit.exe' in '.' or 'kit'
        # Linux: 'kit.sh' in '.' or 'kit' followed by 'kit' in '.' or 'kit'
        exe_names = ["kit.sh", "kit"]
        if sys.platform.startswith("win"):
            exe_names = ["kit.bat", "kit.exe"]

        # look in 4 places...
        for dir_name in [appdir, os.path.join(appdir, "kit")]:
            for exe_name in exe_names:
                if os.path.exists(os.path.join(dir_name, exe_name)):
                    return os.path.join(dir_name, exe_name)

        return None

    def _check_modules(self) -> None:
        """Verify that the Python interpreter is correct

        Check for module dependencies. If not present, raise an exception.

        Raises
        ------
        RuntimeError
            if the necessary modules are missing.

        """
        # One time check for this
        if len(self._interpreter):
            return

        # if a module, then we are inside EnSight
        if isinstance(self._ensight, ModuleType):  # pragma: no cover
            # in this case, we can just use cpython
            import ceiversion
            import enve

            cei_home = os.environ.get("CEI_HOME", enve.home())
            self._interpreter = os.path.join(cei_home, "bin", f"cpython{ceiversion.apex_suffix}")
            if platform.system() == "Windows":
                self._interpreter += ".bat"
            return
        # Using the python interpreter running this code
        self._interpreter = sys.executable

        # in the future, these will be part of the pyensight wheel
        # dependencies, but for now we include this check.
        try:
            import pxr  # noqa: F401
            import pygltflib  # noqa: F401
        except Exception:
            raise RuntimeError("Unable to detect omniverse dependencies: usd-core, pygltflib.")

    def is_running_omniverse(self) -> bool:
        """Check that an Omniverse connection is active

        Returns
        -------
        bool
            True if the connection is active, False otherwise.
        """
        if self._server_pid is None:
            return False
        if psutil.pid_exists(self._server_pid):
            return True
        self._server_pid = None
        return False

    def create_connection(
        self,
        omniverse_path: str,
        include_camera: bool = False,
        normalize_geometry: bool = False,
        temporal: bool = False,
        live: bool = True,
        debug_filename: str = "",
        time_scale: float = 1.0,
        options: dict = {},
    ) -> None:
        """Ensure that an EnSight dsg -> omniverse server is running

        Connect the current EnSight session to an Omniverse server.
        This is done by launching a new service that makes a dynamic scene graph
        connection to the EnSight session and pushes updates to the Omniverse server.
        The initial EnSight scene will be pushed after the connection is established.

        Parameters
        ----------
        omniverse_path : str
            The directory name where the USD files should be saved. For example:
            "C:/Users/test/OV/usdfiles"
        include_camera : bool
            If True, apply the EnSight camera to the Omniverse scene.  This option
            should be used if the target viewer is in AR/VR mode.  Defaults to False.
        normalize_geometry : bool
            Omniverse units are in meters.  If the source dataset is not in the correct
            unit system or is just too large/small, this option will remap the geometry
            to a unit cube.  Defaults to False.
        temporal : bool
            If True, save all timesteps.
        live : bool
            If True, one can call 'update()' to send updated geometry to Omniverse.
            If False, the Omniverse connection will push a single update and then
            disconnect.  Defaults to True.
        time_scale : float
            Multiply all EnSight time values by this factor before exporting to Omniverse.
            The default is 1.0.
        debug_filename : str
            If the name of a file is provided, it will be used to save logging information on
            the connection between EnSight and Omniverse.  This option is no longer supported,
            but the API remains for backwards compatibility.
        options : dict
            Allows for a fallback for the grpc host/port and the security token.
        """
        if not isinstance(self._ensight, ModuleType):
            self._ensight._session.ensight_version_check("2023 R2")
        self._check_modules()
        if self.is_running_omniverse():
            raise RuntimeError("An Omniverse server connection is already active.")
        dsg_uri = None
        is_win = "Win" in platform.system()
        grpc_use_tcp_sockets = False
        grpc_allow_network_connectsion = False
        grpc_disable_tls = False
        if not isinstance(self._ensight, ModuleType):
            # Make sure the internal ui module is loaded
            grpc_use_tcp_sockets = self._ensight._session._grpc_use_tcp_sockets
            grpc_allow_network_connectsion = self._ensight._session._grpc_allow_network_connections
            grpc_disable_tls = self._ensight._session._grpc_disable_tls
            self._ensight._session.cmd("import enspyqtgui_int", do_eval=False)
            # Get the gRPC connection details and use them to launch the service
            use_tcp_sockets = self._ensight._session._grpc_use_tcp_sockets
            token = self._ensight._session.grpc.security_token
            if not is_win and not use_tcp_sockets:
                uds_path = self._ensight._session._grpc_uds_pathname
                dsg_uds_path = "/tmp/greeter"
                if uds_path:
                    dsg_uds_path = uds_path
                dsg_uri = f"unix:{dsg_uds_path}.sock"
            else:
                port = self._ensight._session.grpc.port()
                hostname = self._ensight._session.grpc.host
                token = self._ensight._session.grpc.security_token
                dsg_uri = f"grpc://{hostname}:{port}"
        else:
            hostname = options.get("host", "127.0.0.1")
            port = options.get("port", 12345)
            uds_path = options.get("uds_path")
            token = options.get("security", "")
            if uds_path and not is_win:
                dsg_uri = f"unix:{uds_path}.sock"
            else:
                dsg_uri = f"grpc://{hostname}:{port}"

        # Launch the server via the 'ansys.pyensight.core.utils.omniverse_cli' module
        cmd = [self._interpreter]
        cmd.extend(["-m", "ansys.pyensight.core.utils.omniverse_cli"])
        cmd.append(omniverse_path)
        if token:
            cmd.extend(["--security_token", token])
        if temporal:
            cmd.extend(["--temporal", "true"])
        if not include_camera:
            cmd.extend(["--include_camera", "false"])
        if normalize_geometry:
            cmd.extend(["--normalize_geometry", "true"])
        if time_scale != 1.0:
            cmd.extend(["--time_scale", str(time_scale)])
        if not live:
            cmd.extend(["--oneshot", "1"])
        if grpc_allow_network_connectsion:
            cmd.extend(["--grpc_allow_network_connections", "1"])
        if grpc_disable_tls:
            cmd.extend(["--grpc_disable_tls", "1"])
        if grpc_use_tcp_sockets:
            cmd.extend(["--grpc_use_tcp_sockets", "1"])
        cmd.extend(["--dsg_uri", dsg_uri])
        env_vars = os.environ.copy()
        # we are launching the kit from EnSight or PyEnSight.  In these cases, we
        # inform the kit instance of:
        # (1) the name of the "server status" file, if any
        self._new_status_file()
        env_vars["ANSYS_OV_SERVER_STATUS_FILENAME"] = self._status_filename
        process = subprocess.Popen(cmd, close_fds=True, env=env_vars)
        self._server_pid = process.pid

    def _new_status_file(self, new=True) -> None:
        """
        Remove any existing status file and create a new one if requested.

        Parameters
        ----------
        new : bool
            If True, create a new status file.
        """
        if self._status_filename:
            try:
                os.remove(self._status_filename)
            except OSError:
                pass
        self._status_filename = ""
        if new:
            self._status_filename = os.path.join(
                tempfile.gettempdir(), str(uuid.uuid1()) + "_gs_status.txt"
            )

    def read_status_file(self) -> dict:
        """Read the status file and return its contents as a dictionary.

        Note: this can fail if the file is being written to when this call is made, so expect
        failures.

        Returns
        -------
        Optional[dict]
            A dictionary with the fields 'status', 'start_time', 'processed_buffers', 'total_buffers' or empty
        """
        if not self._status_filename:
            return {}
        try:
            with open(self._status_filename, "r") as status_file:
                data = json.load(status_file)
        except Exception:
            return {}
        return data

    def close_connection(self) -> None:
        """Shut down the open EnSight dsg -> omniverse server

        Break the connection between the EnSight instance and Omniverse.

        """
        self._check_modules()
        if not self.is_running_omniverse():
            return
        proc = psutil.Process(self._server_pid)
        for child in proc.children(recursive=True):
            if psutil.pid_exists(child.pid):
                # This can be a race condition, so it is ok if the child is dead already
                try:
                    child.kill()
                except psutil.NoSuchProcess:
                    pass
        # Same issue, this process might already be shutting down, so NoSuchProcess is ok.
        try:
            proc.kill()
        except psutil.NoSuchProcess:
            pass
        self._server_pid = None
        self._new_status_file(new=False)

    def update(self, temporal: bool = False) -> None:
        """Update the geometry in Omniverse

        Export the current EnSight scene to the current Omniverse connection.

        Parameters
        ----------
        temporal : bool
            If True, export all timesteps.
        """
        update_cmd = "dynamicscenegraph://localhost/client/update"
        if temporal:
            update_cmd += "?timesteps=1"
        self._check_modules()
        if not self.is_running_omniverse():
            raise RuntimeError("No Omniverse server connection is currently active.")
        if not isinstance(self._ensight, ModuleType):
            self._ensight._session.ensight_version_check("2023 R2")
            cmd = f'enspyqtgui_int.dynamic_scene_graph_command("{update_cmd}")'
            self._ensight._session.cmd(cmd, do_eval=False)
        else:
            import enspyqtgui_int

            enspyqtgui_int.dynamic_scene_graph_command(f"{update_cmd}")
