# Copyright (C) 2023 - 2026 ANSYS, Inc. and/or its affiliates.
# SPDX-License-Identifier: MIT
#
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Workbench client module for PyWorkbench."""

import glob
import json
import logging
import logging.handlers
from logging.handlers import WatchedFileHandler
import os
import re
import time

import tqdm

from ansys.api.workbench.v0 import workbench_pb2 as wb
from ansys.api.workbench.v0.workbench_pb2_grpc import WorkbenchServiceStub
from ansys.tools.common.cyberchannel import create_channel
from ansys.workbench.core.example_data import ExampleData


class WorkbenchClient:
    """Functions of a PyWorkbench client.

    Parameters
    ----------
    local_workdir : str
        Local working directory for the client.
    server_host : str
        Hostname or IP address of the server.
    server_port : int
        Port number of the server.
    server_security : string
        Security mode of the server.
        Options are: "insecure", "uds", "wnua", "mtls"
    """

    def __init__(self, local_workdir, server_host, server_port, server_security):
        """Create a Workbench client."""
        self.workdir = local_workdir
        self._server_host = server_host
        self._server_port = server_port
        self._server_security = server_security
        self._server_version = -1
        self.__init_logging()

    @property
    def server_version(self):
        """The Workbench version of the connected server."""
        if self._server_version <= 0:
            ver_str = self.run_script_string("""import json
wb_script_result=json.dumps(GetFrameworkVersion())""")
            ver_str = ver_str.replace(".", "")
            self._server_version = int(ver_str)
        return self._server_version

    def __enter__(self):
        """Connect to the server when entering a context."""
        self._connect()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        """Disconnect from the server when exiting a context."""
        self._disconnect()

    def _connect(self):
        """Connect to the server."""
        self.channel = create_channel(
            host=self._server_host,
            port=self._server_port,
            transport_mode=self._server_security,
            certs_dir=None,
            grpc_options=None,
        )
        self.stub = WorkbenchServiceStub(self.channel)
        logging.info(
            f"connected to the WB server at {self._server_host}:{self._server_port} "
            "using {self._server_security} connection"
        )

    def _disconnect(self):
        """Disconnect from the server."""
        if self.channel:
            self.channel.close()
            self.channel = None
            self.stub = None
            self._server_version = -1
            logging.info("Disconnected from the Workbench server")

    def _is_connected(self):
        """Return whether this client is connected to the server."""
        return self.channel is not None

    def __init_logging(self):
        """Initialize logging for the client."""
        self._logger = logging.getLogger("WB")
        self._logger.setLevel(logging.DEBUG)
        self._logger.propagate = False
        stream_handler = logging.StreamHandler()
        stream_handler.setFormatter(logging.Formatter("%(levelname)s: %(message)s"))
        stream_handler.setLevel(logging.WARNING)
        self._logger.addHandler(stream_handler)
        self.__log_console_handler = stream_handler

    def set_console_log_level(self, log_level):
        """Set the log filter level for the client console.

        Parameters
        ----------
        log_level : str, default: "error"
            Level of logging. Options are "critical", "debug", "error", "info", and "warning".
        """
        self.__log_console_handler.setLevel(WorkbenchClient.__to_python_log_level(log_level))

    def set_log_file(self, log_file):
        """Set a local log file for the Workbench server log.

        Create a local log file if one does not exist and append it to the existing log file.

        Parameters
        ----------
        log_file : str
            Path to a local file to use for logging.
        """
        self.reset_log_file()

        file_handler = WatchedFileHandler(log_file)
        file_handler.setFormatter(logging.Formatter("%(levelname)s: %(message)s"))
        file_handler.setLevel(logging.DEBUG)
        self._logger.addHandler(file_handler)
        self.__log_file_handler = file_handler

    def reset_log_file(self):
        """No longer use the current log file for the Workbench server log."""
        if self.__log_file_handler is None:
            return
        self.__log_file_handler.close()
        self._logger.removeHandler(self.__log_file_handler)
        self.__log_file_handler = None

    __log_file_handler = None
    __log_console_handler = None

    def run_script_string(self, script_string, args=None, log_level="error"):
        """Run a script as given in the input string on the server.

        Parameters
        ----------
        script_string : str
            String containing the content of the script to run.
        args : dictionary of script variable names and values
            Variables in the script specified as $$varname%%50%% will be converted to variable
            values or use the default value - 50 in the example.
        log_level : str, default: "error"
            Level of logging. Options are "critical", "debug", "error", "info", and "warning".

        Returns
        -------
        str
            Output defined in the script.

        Examples
        --------
        Run a Workbench script, given in a string, that returns the name of
        a newly created system.

        >>> wb.run_script_string(r'''import json
        wb_script_result=json.dumps(GetTemplate(TemplateName="FLUENT").CreateSystem().Name)
        ''')

        """
        if not self._is_connected():
            logging.error("Workbench client is not yet connected to a server.")
            return None
        updated_script_string = script_string
        if args and len(args) > 0:
            if any(not re.match(r"^\w+$", arg) for arg in args.keys()):
                logging.error("script argument name contains illegal character.")
                return None
            for arg_name in args:
                updated_script_string = re.sub(
                    r"\$\$" + arg_name + r"%%((?!%%).)*%%",
                    str(args[arg_name]),
                    updated_script_string,
                )
        updated_script_string = re.sub(r"\$\$\w+%%(((?!%%).)*)%%", r"\1", updated_script_string)
        request = wb.RunScriptRequest(
            content=updated_script_string,
            log_level=WorkbenchClient.__to_server_log_level(log_level),
        )
        for response in self.stub.RunScript(request):
            if response.log and response.log.messages and len(response.log.messages) > 0:
                for log_entry in response.log.messages:
                    self.__python_logging(log_entry.level, log_entry.message)
            if response.result:
                if response.result.error:
                    logging.error("Error when running the script: " + response.result.error)
                    return None
                elif response.result.result:
                    logging.info("The script has finisished.")
                    return json.loads(response.result.result)

    def run_script_file(self, script_file_name, args=None, log_level="error"):
        """Run a script file on the server.

        Parameters
        ----------
        script_file_name : str
            Name of the script file to run. The script file should be located in the client
            working directory
        args : dictionary of script variable names and values
            Variables in the script specified as $$varname%%50%% will be converted to variable
            values or use the default value - 50 in the example.
        log_level : str, default: "error"
            Level of logging. Options are "critical", "debug", "error", "info", and "warning".

        Returns
        -------
        str
            Output defined in the script.
        """
        if not self._is_connected():
            logging.error("Workbench client is not yet connected to a server")
            return None
        script_path = os.path.join(self.workdir, script_file_name)
        with open(script_path, encoding="utf-8-sig") as sf:
            script_string = sf.read()
        return self.run_script_string(script_string, args, log_level)

    def upload_file(self, *file_list, show_progress=True):
        """Upload one or more files from the client to the server.

        Parameters
        ----------
        file_list : list[str]
            List of paths to the one or more local files to upload. The
            wildcard characters "?" and "*" are supported.
        show_progress : bool, default: True
            Whether to show a progress bar during the upload.

        Returns
        -------
        list[str]
            Names of the uploaded files.
        """
        if not self._is_connected():
            logging.error("Workbench client is not yet connected to a server.")
            return
        requested = []
        for file_pattern in file_list:
            if "*" in file_pattern or "?" in file_pattern:
                if not os.path.isabs(file_pattern):
                    file_pattern = os.path.join(self.workdir, file_pattern)
                requested.extend(glob.glob(file_pattern))
            else:
                requested.append(file_pattern)
        existing_files = []
        nonexisting_files = []
        for file_name in requested:
            if not os.path.isabs(file_name):
                file_name = os.path.join(self.workdir, file_name)
            if os.path.isfile(file_name):
                existing_files.append(file_name)
            else:
                nonexisting_files.append(file_name)
        if len(nonexisting_files) > 0:
            logging.warning(
                "The following files do not exist and are skipped: " + "\n".join(nonexisting_files)
            )
        for file_path in existing_files:
            logging.info(f"uploading file {file_path}")
            response = self.stub.UploadFile(self.__upload_iterator(file_path, show_progress))
            if response.error:
                logging.error("Error during file upload: " + response.error)
            else:
                logging.info(
                    "A file is uploaded to the server with the name: " + response.file_name
                )

    def __upload_iterator(self, file_path, show_progress):
        dir_path, file_name = os.path.split(file_path)
        yield wb.UploadFileRequest(file_name=file_name)

        pbar = None
        if show_progress:
            bytes = os.path.getsize(file_path)
            pbar = tqdm.tqdm(
                total=bytes,
                desc=f"Uploading {file_name}",
                unit="B",
                unit_scale=True,
                unit_divisor=1024,
            )

        chunk_size = 65536  ## 64 kb
        with open(file_path, mode="rb") as f:
            while True:
                chunk = f.read(chunk_size)
                if chunk:
                    if pbar is not None:
                        pbar.update(len(chunk))
                    yield wb.UploadFileRequest(file_content=chunk)
                else:
                    if pbar is not None:
                        pbar.close()
                    return

    def upload_file_from_example_repo(self, relative_file_path, show_progress=True):
        """Upload a file from the Ansys ``example-data`` repository to the server.

        Parameters
        ----------
        relative_file_path : str
            File path relative to the ``pyworkbench`` folder in the ``example-data`` repository.
        show_progress : bool, default: True
            Whether to show a progress bar during the upload.
        """
        if not self._is_connected():
            logging.error("Workbench client is not yet connected to a server.")
            return
        downloaded = ExampleData.download(relative_file_path, self.workdir)
        self.upload_file(downloaded, show_progress=show_progress)

    def download_file(self, file_name, show_progress=True, target_dir=None):
        """Download one or more files from the server.

        Parameters
        ----------
        file_name : str
            Name of the file. File must be located in the server's working directory.
            The wildcard characters "?" and "*" are supported. A ZIP file is automatically
            generated and downloaded when multiple files are specified.
        show_progress : bool, default: True
            Whether to show a progress bar during the download.
        target_dir : str, default: None
            Path to a local directory to download the files to. The default is ``None``,
            in which case the client working directory is used.


        Returns
        -------
        str
            Name of the downloaded file.
        """
        if not self._is_connected():
            logging.error("Workbench client is not yet connected to a server.")
            return None
        request = wb.DownloadFileRequest(file_name=file_name)
        file_name = file_name.replace("*", "_").replace("?", "_")
        td = target_dir
        if td is None:
            td = self.workdir
        file_path = os.path.join(td, file_name)
        pbar = None
        started = False
        for response in self.stub.DownloadFile(request):
            if response.error:
                logging.error("Error during file download: " + response.error)
                return None
            if response.file_info:
                if response.file_info.is_archive:
                    file_name += ".zip"
                    file_path += ".zip"
                if response.file_info.file_size > 0:
                    if show_progress:
                        pbar = tqdm.tqdm(
                            total=response.file_info.file_size,
                            desc=f"Downloading {file_name}",
                            unit="B",
                            unit_scale=True,
                            unit_divisor=1024,
                        )
            if response.file_content:
                size = len(response.file_content)
                if size > 0:
                    if not started:
                        if os.path.exists(file_path):
                            os.remove(file_path)
                        started = True

                    try_more = 3
                    while try_more > 0:
                        try:
                            with open(file_path, mode="ab") as f:
                                f.write(response.file_content)
                            try_more = 0
                        except PermissionError:
                            # intermittent "Permission Denied" error can happen
                            if try_more <= 0:
                                raise
                            try_more -= 1
                            time.sleep(0.1)

                    if pbar is not None:
                        pbar.update(size)
        logging.info(f"Dwnloaded the file {file_name}.")
        if pbar is not None:
            pbar.close()
        return file_name

    def download_project_archive(
        self, archive_name, include_solution_result_files=True, show_progress=True
    ):
        """Create and download the project archive.

        Parameters
        ----------
        archive_name : str
            Name of the project archive to use, without the file extension.
        include_solution_result_files : bool, default: True
            Whether to include solution and result files in the archive.
        show_progress : bool, default: True
            Whether to show a progress bar during the download.
        """
        if not re.match(r"^\w+$", archive_name):
            logging.error("archive name should contain only alphanumeric characters")
            return
        script = f"""import os
import json
successful = False
wd = GetServerWorkingDirectory()
if os.path.basename(GetProjectFile()).StartsWith("wbnew."):
    Save(FilePath=os.path.join(wd, "{archive_name}.wbpj"), Overwrite=True)
else:
    Save(Overwrite=True)
Archive(FilePath=os.path.join(wd, "{archive_name}.wbpz"),
    IncludeSkippedFiles={include_solution_result_files})
successful = True
wb_script_result =json.dumps(successful)
"""
        archive_created = self.run_script_string(script)
        if not archive_created:
            logging.error(
                (
                    "Failed to create the project archive. "
                    "Make sure that the solver PyAnsys sessions are closed."
                )
            )
            return
        self.download_file(archive_name + ".wbpz", show_progress=show_progress)

    def __python_logging(self, log_level, msg):
        """Log a message with the given log level.

        Parameters
        ----------
        log_level : int
            log level. Options are ``logging.CRITICAL``, ``logging.DEBUG``, ``logging.ERROR``,
            ``logging.INFO``, and ``logging.WARNING``.

        msg : str
            Message to log.
        """
        if log_level == wb.LOG_DEBUG:
            self._logger.debug(msg)
        elif log_level == wb.LOG_INFO:
            self._logger.info(msg)
        elif log_level == wb.LOG_WARNING:
            self._logger.warn(msg)
        elif log_level == wb.LOG_ERROR:
            self._logger.error(msg)
        elif log_level == wb.LOG_FATAL:
            self._logger.fatal(msg)

    @staticmethod
    def __to_python_log_level(log_level):
        """Convert the given log level to the corresponding Python log level.

        Parameters
        ----------
        log_level : str
            level of logging: options are "debug", "info", "warning", "error", "critical"

        Returns
        -------
        int
            the corresponding Python log level
        """
        log_level = log_level.lower()
        for level_name, server_level in WorkbenchClient.__log_levels.items():
            if log_level in level_name:
                return server_level[1]
        return logging.NOTSET

    @staticmethod
    def __to_server_log_level(log_level):
        """Convert the given log level to the corresponding server log level.

        Parameters
        ----------
        log_level : str
            Level of logging. Options are "critical", "debug", "error", "info", and "warning".

        Returns
        -------
        int
            Corresponding server log level.
        """
        log_level = log_level.lower()
        for level_name, server_level in WorkbenchClient.__log_levels.items():
            if log_level in level_name:
                return server_level[0]
        return wb.LOG_NONE

    __log_levels = {
        "none null": (wb.LOG_NONE, logging.NOTSET),
        "debug": (wb.LOG_DEBUG, logging.DEBUG),
        "information": (wb.LOG_INFO, logging.INFO),
        "warning": (wb.LOG_WARNING, logging.WARNING),
        "error": (wb.LOG_ERROR, logging.ERROR),
        "fatal critical": (wb.LOG_FATAL, logging.CRITICAL),
    }

    def start_mechanical_server(self, system_name, port=0):
        """Start the PyMechanical server for the given system in the Workbench project.

        Parameters
        ----------
        system_name : str
            Name of the system in the Workbench project.
        port : int, default: 0
            Port to use for mechanical server if possible

        Returns
        -------
        int
            Port of the PyMechanical server to use to start the PyMechanical client.

        Examples
        --------
        Start a PyMechanical session for the given system.

        >>> from ansys.mechanical.core import connect_to_mechanical
        >>> server_port=wb.start_mechanical_server(system_name=mech_system_name)
        >>> mechanical = connect_to_mechanical(port=server_port)

        """
        port_arg = ""
        if port > 0:
            port_arg = ", Port=" + str(port)
        pymech_port = self.run_script_string(
            f"""import json
if float(GetFrameworkVersion()) >= 25.2:
    server_port=LaunchMechanicalServerOnSystem(SystemName="{system_name}{port_arg}")
else:
    server_port=LaunchMechanicalServerOnSystem(SystemName="{system_name}")
wb_script_result=json.dumps(server_port)
"""
        )
        return pymech_port

    def stop_mechanical_server(self, system_name):
        """Stop the PyMechanical server for the given system in the Workbench project.

        Parameters
        ----------
        system_name : str
            Name of the system in the Workbench project.

        Examples
        --------
        Stop the PyMechanical session for the given system.

        >>> wb.stop_mechanical_server(system_name=mech_system_name)

        """
        self.run_script_string(f"""if float(GetFrameworkVersion()) >= 25.2:
    StopMechanicalServerOnSystem(SystemName="{system_name}")
""")

    def start_fluent_server(self, system_name):
        """Start the PyFluent server for the given system in the Workbench project.

        Parameters
        ----------
        system_name : str
            Name of the system in the Workbench project.

        Returns
        -------
        str
            Path to the local file containing the PyFluent server information that
            can be used to start a PyFluent client.

        Examples
        --------
        Start a PyFluent session for the given system.

        >>> import ansys.fluent.core as pyfluent
        >>> server_info_file=wb.start_fluent_server(system_name=fluent_sys_name)
        >>> fluent=pyfluent.connect_to_fluent(server_info_file_name=server_info_file)

        """
        server_info_file_name = self.run_script_string(
            f"""import json
server_info_file=LaunchFluentServerOnSystem(SystemName="{system_name}")
wb_script_result=json.dumps(server_info_file)
"""
        )
        local_copy = os.path.join(self.workdir, server_info_file_name)
        if os.path.exists(local_copy):
            os.remove(local_copy)
        self.download_file(server_info_file_name, show_progress=False)
        return local_copy

    def stop_fluent_server(self, system_name):
        """Stop the Fluent server for the given system in the Workbench project.

        Parameters
        ----------
        system_name : str
            Name of the system in the Workbench project.

        Examples
        --------
        Stop the Fluent session for the given system.

        >>> wb.stop_fluent_server(system_name=mech_system_name)

        """
        self.run_script_string(f"""if float(GetFrameworkVersion()) >= 25.2:
    StopFluentServerOnSystem(SystemName="{system_name}")
""")

    def start_sherlock_server(self, system_name):
        """Start the PySherlock server for the given system in the Workbench project.

        Parameters
        ----------
        system_name : str
            Name of the system in the Workbench project.

        Returns
        -------
        int
            Port of the PySherlock server to use to start a PySherlock client.

        Examples
        --------
        Start a PySherlock session for the given system.

        >>> from ansys.sherlock.core import pysherlock
        >>> server_port=wb.start_sherlock_server(system_name=sherlock_system_name)
        >>> sherlock = pysherlock.connect_grpc_channel(port=server_port)
        >>> sherlock.common.check()

        """
        pysherlock_port = self.run_script_string(
            f"""import json
server_port=LaunchSherlockServerOnSystem(SystemName="{system_name}")
wb_script_result=json.dumps(server_port)
"""
        )
        return pysherlock_port

    def stop_sherlock_server(self, system_name):
        """Stop the Sherlock server for the given system in the Workbench project.

        Parameters
        ----------
        system_name : str
            Name of the system in the Workbench project.

        Examples
        --------
        Stop the Sherlock session for the given system.

        >>> wb.stop_sherlock_server(system_name=mech_system_name)

        """
        self.run_script_string(f"""if float(GetFrameworkVersion()) >= 25.2:
    StopSherlockServerOnSystem(SystemName="{system_name}")
""")


__all__ = ["WorkbenchClient"]
