# SPDX-FileCopyrightText: Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.

import os
import socket
import struct
import sys
import time
import subprocess

import NsysServer_pb2


class NsysClient:
    """
        A class that implements methods for the Nsys client.

        The client uses a Unix socket to connect with the server.
    """

    def __init__(self, socket_name="nsys_socket", launch_server=False, connect_now=True):
        self.socket_name = socket_name
        self.client_socket = None
        self.report_path = None
        self.server_proc = None
        src_path=os.path.realpath(os.path.dirname(__file__))
        self.server_install_path=os.path.realpath(src_path+"/../../../")
        if launch_server:
            self.launchServer(socket_name, self.server_install_path)
            time.sleep(1)
        if connect_now:
            self.connect_to_socket()

    def __del__(self):
        self.exit()

    def launchServer(self, socket_name, server_install_path):
        # socket_name is the name of the server socket to use for the server
        # server_path is the nsys installation path to use to find the CLI (not including 'bin/nsys') 
        proc = subprocess.Popen([server_install_path+"/bin/nsys","server","--socketName",socket_name])
        self.server_proc = proc

    def print_usage(self):
        """
            A function that prints a help menu, clarifying the usage of the NsysClient class.
        """
        print(
            """
    This class allows a user to request reports and other information from
    Nsight Systems. It uses a UNIX socket to establish a connection with the
    Nsight Systems GUI.
    An object of this class will connect to the Nsight server and wait for a command
    from the user. Please type a command when prompted.
    Available commands are:
        - list-reports : The command lists the available report files on the directory that
                         the Nsight Systems server was launched from.
        - get-report : The command requests access to a report and receives the path to an
                       arrow IPC format report.
        - exit : The command closes the socket from the client side.
            """
        )

    def connect_to_socket(self):
        """
            A function that tries to establish a Unix socket connection with the Nsight Systems
            server.
        """
        socket_path = "/tmp/"+self.socket_name
        if os.path.exists(socket_path):
            self.client_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            self.client_socket.connect(socket_path)
            #print("Client connected.")
        else:
            print("Client couldn't connect to NsysServer!")

    def read_and_send_command(self):
        """
            A function that prompts the user for a command and sends the command to the server.

            If the socket connection is not established before the function is called, the function
            first attempts to connect to the socket.
        """
        if self.client_socket is None:
            self.connect_to_socket()

        while True:
            try:
                command = input("> ")
                if not self.send_command(command):
                    break
            except KeyboardInterrupt:
                print("Shutting down.")

        self.client_socket.close()

    def send_command(self, command):
        """
            A function that calls the appropriate command-function, based on the command typed by the user.
        """
        if command != "":
            if command == "list-reports":
                self.list_reports()
            elif command == "get-report":
                print("Please provide the name of the report to receive.")
                report_path = input("> ")
                self.get_report(report_path)
            elif command == "exit":
                self.exit()
                return False
            elif command == "shutdown":
                self.exit(shutdown=True)
            else:
                print("Unrecognized command provided")
        return True

    def exit(self, shutdown=False):
        """
            A function that closes the socket from the client side.
        """
        if not self.client_socket:
            self.connect_to_socket()
        if self.client_socket:
            if shutdown:
                command = "shutdown"
            else:
                command = "exit"
            exit_command = NsysServer_pb2.clientCommands()
            exit_command.command = NsysServer_pb2.clientCommands.CommandType.Value(
                command
            )
            sizeOfCommand = exit_command.ByteSize()
            self.client_socket.sendall(sizeOfCommand.to_bytes(4, byteorder="big"))
            self.client_socket.sendall(exit_command.SerializeToString())
            # Close the socket
            self.client_socket.close()
        self.client_socket = None
        if self.server_proc:
            self.server_proc.terminate()
            self.server_procs = None

    def list_reports(self):
        """
            A function that implements the list-reports command.

            It receives a list of available reports from the server
        """

        def print_report_list(report_list):
            """
                A function that supports the list-reports command.

                This function receives a list of strings.
                It prints each report name on a new line
            """
            for report in report_list:
                print(report)

        list_reports_command = NsysServer_pb2.clientCommands()
        list_reports_command.command = NsysServer_pb2.clientCommands.CommandType.Value(
            "listReports"
        )
        sizeOfCommand = list_reports_command.ByteSize()
        self.client_socket.sendall(sizeOfCommand.to_bytes(4, byteorder="big"))
        self.client_socket.sendall(list_reports_command.SerializeToString())

        print("Receiving data...")

        # Receive path to requested report
        unsigned_long_size = 8
        data_size_bytes = self.client_socket.recv(
            unsigned_long_size, socket.MSG_WAITALL
        )

        # Little-endian interpretation
        data_size = int.from_bytes(data_size_bytes, "little")
        data = self.client_socket.recv(data_size, socket.MSG_WAITALL)
        reports_list = NsysServer_pb2.listReportsResponse()
        reports_list.ParseFromString(data)
        print_report_list(reports_list.reportFile)

    def get_report(self, report_name, export_type="Arrow"):
        """
            A function that implements the get-report command.

            It sends a report name to the server. If the report exists, the server shares the path
            to an arrow IPC report.
        """
        get_report_command = NsysServer_pb2.clientCommands()
        get_report_command.command = NsysServer_pb2.clientCommands.CommandType.Value(
            "getReport"
        )
        sizeOfCommand = get_report_command.ByteSize()
        self.client_socket.sendall(sizeOfCommand.to_bytes(4, byteorder="big"))
        self.client_socket.sendall(get_report_command.SerializeToString())

        # Send requested report name
        get_report_request = NsysServer_pb2.getReportRequest()
        get_report_request.reportName = report_name
        get_report_request.exportType = export_type
        sizeOfReportRequest = get_report_request.ByteSize()
        self.client_socket.sendall(sizeOfReportRequest.to_bytes(4, byteorder="big"))
        self.client_socket.sendall(get_report_request.SerializeToString())

        # Receive path to requested report
        unsigned_long_size = 8
        data_size_bytes = self.client_socket.recv(
            unsigned_long_size, socket.MSG_WAITALL
        )

        # Little-endian interpretation
        data_size = int.from_bytes(data_size_bytes, "little")
        data = self.client_socket.recv(data_size, socket.MSG_WAITALL)
        get_report_response = NsysServer_pb2.getReportResponse()
        get_report_response.ParseFromString(data)
        self.report_path = get_report_response.pathToReport
        if get_report_response.requestedReportExists:
            print("Report path is:", self.report_path)
        else:
            print("Report does not exist.")
        return self.report_path

    def open_report_ui(self, report_name):
        """
            open a report on the UI

            It sends a report name to the server. If the report exists, the server shares the path
            to an arrow IPC report.
        """
        open_report_command = NsysServer_pb2.clientCommands()
        open_report_command.command = NsysServer_pb2.clientCommands.CommandType.Value(
            "openReportUI"
        )
        sizeOfCommand = open_report_command.ByteSize()
        self.client_socket.sendall(sizeOfCommand.to_bytes(4, byteorder="big"))
        self.client_socket.sendall(open_report_command.SerializeToString())

        # Send requested report name
        get_report_request = NsysServer_pb2.getReportRequest()
        get_report_request.reportName = report_name
        sizeOfReportRequest = get_report_request.ByteSize()
        self.client_socket.sendall(sizeOfReportRequest.to_bytes(4, byteorder="big"))
        self.client_socket.sendall(get_report_request.SerializeToString())


def main():
    """
        A function to drive and check the NsysClient class
    """
    nsys_client = NsysClient()

    nsys_client.print_usage()

    nsys_client.list_reports()

    report_name = "this_report_doesnt_exist"
    nsys_client.get_report(report_name)

    nsys_client.exit()


if __name__ == "__main__":
    main()