# This file is part of Xpra.
# Copyright (C) 2010-2024 Antoine Martin <antoine@xpra.org>
# Copyright (C) 2008, 2010 Nathaniel Smith <njs@pobox.com>
# Xpra is released under the terms of the GNU GPL v2, or, at your option, any
# later version. See the file COPYING for details.

import os.path
import sys
from time import monotonic
from typing import Any

from xpra.util.objects import typedict
from xpra.util.str_fn import nonl, csv, ellipsizer, repr_ellipsized, sorted_nicely, bytestostr, hexstr
from xpra.util.env import envint, first_time
from xpra.common import ConnectionMessage, disconnect_is_an_error
from xpra.os_util import gi_import, get_hex_uuid, POSIX, OSX
from xpra.util.io import stderr_print
from xpra.net.common import PacketType
from xpra.util.stats import std_unit
from xpra.client.base.client import EXTRA_TIMEOUT
from xpra.client.base.gobject import GObjectXpraClient
from xpra.exit_codes import ExitCode, ExitValue
from xpra.log import Logger

log = Logger("gobject", "client")

GLib = gi_import("GLib")


def errwrite(msg) -> None:
    stderr_print(msg)


class CommandConnectClient(GObjectXpraClient):
    """
        Utility superclass for clients that only send one command
        via the hello packet.
    """

    def __init__(self, opts):
        super().__init__()
        super().init(opts)
        self.display_desc = {}
        # not used by command line clients,
        # so don't try probing for printers, etc
        self.file_transfer = False
        self.printing = False
        self.command_timeout = None
        # don't bother with many of these things for one-off commands:
        for x in ("ui_client", "windows", "webcam", "keyboard", "mouse", "network-state"):
            self.hello_extra[x] = False
        # for newer versions, it is easier:
        self.hello_extra["wants"] = []

    def setup_connection(self, conn):
        protocol = super().setup_connection(conn)
        if conn.timeout > 0:
            self.command_timeout = GLib.timeout_add((conn.timeout + self.COMMAND_TIMEOUT) * 1000, self.timeout)
        return protocol

    def timeout(self, *_args) -> None:
        log.warn("timeout!")  # pragma: no cover

    def cancel_command_timeout(self) -> None:
        ct = self.command_timeout
        if ct:
            self.command_timeout = None
            GLib.source_remove(ct)

    def _process_connection_lost(self, packet: PacketType) -> None:
        log("_process_connection_lost%s", packet)
        # override so we don't log a warning
        # "command clients" are meant to exit quickly by losing the connection
        p = self._protocol
        if p and p.input_packetcount == 0:
            # never got any data back, so we have never connected,
            # try to give feedback to the user as to why that is:
            details = ""
            if len(packet) > 1:
                details = ": %s" % csv(packet[1:])
            log.warn("Connection failed%s", details)
            self.quit(ExitCode.CONNECTION_FAILED)
        else:
            self.quit(ExitCode.OK)

    def server_connection_established(self, caps: typedict) -> bool:
        # don't bother parsing the network caps:
        # * it could cause errors if the caps are missing
        # * we don't care about sending anything back after hello
        log("server_capabilities: %s", ellipsizer(caps))
        log("protocol state: %s", self._protocol.save_state())
        self.cancel_command_timeout()
        self.do_command(caps)
        return True

    def do_command(self, caps: typedict) -> None:
        raise NotImplementedError()


class SendCommandConnectClient(CommandConnectClient):
    """
        Utility superclass for clients that only send at least one more packet
        after the hello packet.
        So unlike CommandConnectClient, we do need the network and encryption to be setup.
    """

    def server_connection_established(self, caps) -> bool:
        assert self.parse_encryption_capabilities(caps), "encryption failure"
        assert self.parse_network_capabilities(caps), "network capabilities failure"
        return super().server_connection_established(caps)

    def do_command(self, caps: typedict) -> None:
        raise NotImplementedError()


class HelloRequestClient(SendCommandConnectClient):
    """
        Utility superclass for clients that send a server request
        as part of the hello packet.
    """

    def make_hello_base(self) -> dict[str, Any]:
        caps = super().make_hello_base()
        caps.update(self.hello_request())
        return caps

    def timeout(self, *_args) -> None:
        self.warn_and_quit(ExitCode.TIMEOUT, "timeout: server did not disconnect us")

    def hello_request(self):  # pragma: no cover
        raise NotImplementedError()

    def do_command(self, caps: typedict) -> None:
        self.quit(ExitCode.OK)

    def _process_disconnect(self, packet: PacketType) -> None:
        # overridden method, so we can avoid printing a warning,
        # we haven't received the hello back from the server
        # but that's fine for a request client
        info = tuple(nonl(x) for x in packet[1:])
        reason = info[0]
        if disconnect_is_an_error(reason):
            self.server_disconnect_warning(*info)
        elif self.exit_code is None:
            # we're not in the process of exiting already,
            # tell the user why the server is disconnecting us
            self.server_disconnect(*info)


class ScreenshotXpraClient(CommandConnectClient):
    """ This client does one thing only:
        it sends the hello packet with a screenshot request
        and exits when the resulting image is received (or timedout)
    """

    def __init__(self, opts, screenshot_filename):
        self.screenshot_filename = screenshot_filename
        super().__init__(opts)
        self.hello_extra["request"] = "screenshot"

    def timeout(self, *_args) -> None:
        self.warn_and_quit(ExitCode.TIMEOUT, "timeout: did not receive the screenshot")

    def _process_screenshot(self, packet: PacketType) -> None:
        (w, h, encoding, _, img_data) = packet[1:6]
        assert encoding == "png", "expected png screenshot data but got %s" % encoding
        if not img_data:
            self.warn_and_quit(ExitCode.OK,
                               "screenshot is empty and has not been saved"
                               " (maybe there are no windows or they are not currently shown)")
            return
        if self.screenshot_filename == "-":
            output = os.fdopen(sys.stdout.fileno(), "wb", closefd=False)
        else:
            output = open(self.screenshot_filename, "wb")
        with output:
            output.write(img_data)
            output.flush()
        self.warn_and_quit(ExitCode.OK, f"screenshot {w}x{h} saved to: {self.screenshot_filename!r}")

    def init_packet_handlers(self) -> None:
        super().init_packet_handlers()
        self._ui_packet_handlers["screenshot"] = self._process_screenshot


class InfoXpraClient(CommandConnectClient):
    """ This client does one thing only:
        it queries the server with an 'info' request
    """

    def __init__(self, opts):
        super().__init__(opts)
        self.hello_extra["request"] = "info"

    def timeout(self, *_args):
        self.warn_and_quit(ExitCode.TIMEOUT, "timeout: did not receive the info")

    def do_command(self, caps: typedict) -> None:
        if not caps:
            self.quit(ExitCode.NO_DATA)
            return

        def print_fn(s) -> None:
            sys.stdout.write(f"{s}\n")

        def prettify(k, v) -> str:
            ks = str(k)
            if ks.endswith("xid"):
                try:
                    return hex(int(v))
                except (TypeError, ValueError):
                    pass
            if ks.endswith("version") and isinstance(v, (tuple, list)):
                return ".".join(str(x) for x in v)
            if isinstance(v, bytes):
                if ks.endswith(".data"):
                    return hexstr(v)
                return bytestostr(v)
            elif isinstance(v, (tuple, list)):
                eltypes = set(type(x) for x in v)
                if len(eltypes) == 1 and tuple(eltypes)[0] in (int, bool, float):
                    return str(v)
                return str(type(v)(prettify(k, x) for x in v))
            return str(v)

        def print_dict(d: dict, path=""):
            for k in sorted_nicely(d.keys()):
                v = d.get(k)
                fk = prettify(k, k)
                kpath = f"{path}.{fk}".strip(".")
                if isinstance(v, dict):
                    print_dict(v, kpath)
                    continue
                fv = prettify(k, v)
                print_fn(f"{kpath}={nonl(fv)}")

        exit_code = ExitCode.OK
        try:
            print_dict(dict(caps))
        except OSError:
            exit_code = ExitCode.IO_ERROR
        self.quit(exit_code)


class IDXpraClient(InfoXpraClient):

    def __init__(self, *args):
        super().__init__(*args)
        self.hello_extra["request"] = "id"


class RequestXpraClient(CommandConnectClient):

    def __init__(self, request, opts):
        super().__init__(opts)
        self.hello_extra["request"] = request

    def do_command(self, caps: typedict) -> None:
        self.quit(ExitCode.OK)


class ConnectTestXpraClient(CommandConnectClient):
    """ This client does one thing only:
        it queries the server with an 'info' request
    """

    def __init__(self, opts, **kwargs):
        super().__init__(opts)
        self.value = get_hex_uuid()
        self.hello_extra |= {
            "connect_test_request": self.value,
            "request": "connect_test",
            # tells proxy servers we don't want to connect to the real / new instance:
            "connect": False,
        }
        self.hello_extra.update(kwargs)

    def timeout(self, *_args) -> None:
        self.warn_and_quit(ExitCode.TIMEOUT, "timeout: no server response")

    def _process_connection_lost(self, _packet: PacketType) -> None:
        # we should always receive a hello back and call do_command,
        # which sets the correct exit code, landing here is an error:
        self.quit(ExitCode.FAILURE)

    def do_command(self, caps: typedict) -> None:
        if caps:
            ctr = caps.strget("connect_test_response")
            log("do_command(..) expected connect test response='%s', got '%s'", self.value, ctr)
            if ctr == self.value:
                self.quit(ExitCode.OK)
            else:
                self.quit(ExitCode.INTERNAL_ERROR)
        else:
            self.quit(ExitCode.FAILURE)


class MonitorXpraClient(SendCommandConnectClient):
    """ This client does one thing only:
        it prints out events received from the server.
        If the server does not support this feature it exits with an error.
    """

    def __init__(self, opts):
        super().__init__(opts)
        for x in ("features", "events", "request"):
            self.hello_extra.setdefault("wants", []).append(x)
        self.hello_extra["request"] = "event"

    def timeout(self, *args) -> None:
        pass
        # self.warn_and_quit(ExitCode.TIMEOUT, "timeout: did not receive the info")

    def do_command(self, caps: typedict) -> None:
        log.info("waiting for server events")

    def _process_server_event(self, packet: PacketType) -> None:
        log.info(": ".join(str(x) for x in packet[1:]))

    def init_packet_handlers(self) -> None:
        super().init_packet_handlers()
        self._packet_handlers["server-event"] = self._process_server_event
        self._packet_handlers["ping"] = self._process_ping

    def _process_ping(self, packet: PacketType) -> None:
        echotime = packet[1]
        self.send("ping_echo", echotime, 0, 0, 0, -1)


class InfoTimerClient(MonitorXpraClient):
    """
        This client keeps monitoring the server
        and requesting info data
    """
    REFRESH_RATE = envint("XPRA_REFRESH_RATE", 1)

    def __init__(self, *args):
        super().__init__(*args)
        self.info_request_pending = False
        self.server_last_info = typedict()
        self.server_last_info_time = 0
        self.info_timer = 0

    def run(self) -> ExitValue:
        from xpra.gtk.signals import register_os_signals
        register_os_signals(self.signal_handler, "")
        v = super().run()
        self.log("run()=%s" % v)
        self.cleanup()
        return v

    def signal_handler(self, signum: int, *args) -> None:
        self.log(f"exit_code={self.exit_code}")
        self.log(f"signal_handler({signum}, {args})")
        self.quit(128 + signum)
        self.log(f"exit_code={self.exit_code}")

    def log(self, message) -> None:
        # this method is overridden in top client to use a log file
        log(message)

    def err(self, e) -> None:
        log.error(str(e))

    def cleanup(self) -> None:
        self.cancel_info_timer()
        MonitorXpraClient.cleanup(self)

    def do_command(self, caps: typedict) -> None:
        self.send_info_request()
        self.timeout_add(self.REFRESH_RATE * 1000, self.send_info_request)

    def send_info_request(self, *categories) -> bool:
        self.log("send_info_request%s" % (categories,))
        if not self.info_request_pending:
            self.info_request_pending = True
            window_ids = ()  # no longer used or supported by servers
            self.send("info-request", [self.uuid], window_ids, categories)
        if not self.info_timer:
            self.info_timer = self.timeout_add((self.REFRESH_RATE + 2) * 1000, self.info_timeout)
        return True

    def init_packet_handlers(self) -> None:
        MonitorXpraClient.init_packet_handlers(self)
        self.add_packet_handler("info-response", self._process_info_response, False)

    def _process_server_event(self, packet: PacketType) -> None:
        self.log("server event: %s" % (packet,))
        self.last_server_event = packet[1:]
        self.update_screen()

    def _process_info_response(self, packet: PacketType) -> None:
        self.log("info response: %s" % repr_ellipsized(packet))
        self.cancel_info_timer()
        self.info_request_pending = False
        self.server_last_info = typedict(packet[1])
        self.server_last_info_time = monotonic()
        # log.info("server_last_info=%s", self.server_last_info)
        self.update_screen()

    def cancel_info_timer(self) -> None:
        it = self.info_timer
        if it:
            self.info_timer = 0
            self.source_remove(it)

    def info_timeout(self) -> bool:
        self.log("info timeout")
        self.update_screen()
        return True

    def update_screen(self):
        raise NotImplementedError()


def print_prompt() -> None:
    sys.stdout.write("> ")
    sys.stdout.flush()


class ShellXpraClient(SendCommandConnectClient):
    """
        Provides an interactive shell with the socket it connects to
    """

    def __init__(self, opts):
        super().__init__(opts)
        self.stdin_io_watch = None
        self.stdin_buffer = ""
        self.hello_extra["shell"] = "True"

    def timeout(self, *args) -> None:
        """
        The shell client never times out,
        but the superclass calls this method automatically,
        just ignore it.
        """

    def cleanup(self) -> None:
        siw = self.stdin_io_watch
        if siw:
            self.stdin_io_watch = None
            self.source_remove(siw)
        super().cleanup()

    def do_command(self, caps: typedict) -> None:
        if not caps.boolget("shell"):
            msg = "this server does not support the 'shell' subcommand"
            log.error(msg)
            self.disconnect_and_quit(ExitCode.UNSUPPORTED, msg)
            return
        # start reading from stdin:
        self.install_signal_handlers()
        stdin = sys.stdin
        fileno = stdin.fileno()
        import fcntl
        fl = fcntl.fcntl(fileno, fcntl.F_GETFL)
        fcntl.fcntl(fileno, fcntl.F_SETFL, fl | os.O_NONBLOCK)
        self.stdin_io_watch = GLib.io_add_watch(sys.stdin,
                                                GLib.PRIORITY_DEFAULT, GLib.IO_IN,
                                                self.stdin_ready)
        print_prompt()

    def stdin_ready(self, *_args) -> bool:
        data = sys.stdin.read()
        # log.warn("stdin=%r", data)
        self.stdin_buffer += data
        sent = 0
        if self.stdin_buffer.endswith("\n"):
            for line in self.stdin_buffer.splitlines():
                if line:
                    if line.rstrip("\n\r") in ("quit", "exit"):
                        self.disconnect_and_quit(ExitCode.OK, "user requested %s" % line)
                        self.stdin_io_watch = None
                        return False
                    self.send("shell-exec", line.encode())
                    sent += 1
        self.stdin_buffer = ""
        if not sent:
            print_prompt()
        return True

    def init_packet_handlers(self) -> None:
        super().init_packet_handlers()
        self._packet_handlers["shell-reply"] = self._process_shell_reply
        self._packet_handlers["ping"] = self._process_ping

    def _process_ping(self, packet: PacketType) -> None:
        echotime = packet[1]
        self.send("ping_echo", echotime, 0, 0, 0, -1)

    def _process_shell_reply(self, packet: PacketType) -> None:
        fd = int(packet[1])
        message = packet[2]
        if fd == 1:
            stream = sys.stdout
        elif fd == 2:
            stream = sys.stderr
        else:
            raise ValueError(f"invalid file descriptor f{fd}")
        s = str(message)
        if s.endswith("\n"):
            s = s[:-1]
        stream.write("%s" % s)
        stream.flush()
        if fd == 2:
            stream.write("\n")
            print_prompt()


class VersionXpraClient(HelloRequestClient):
    """ This client does one thing only:
        it queries the server for version information and prints it out
    """

    def hello_request(self) -> dict[str, Any]:
        return {
            "request": "version",
            "full-version-request": True,
        }

    def parse_network_capabilities(self, *_args) -> bool:
        # don't bother checking anything - this could generate warnings
        return True

    def do_command(self, caps: typedict) -> None:
        v = caps.strget("version")
        if not v:
            self.warn_and_quit(ExitCode.FAILURE, "server did not provide the version information")
        else:
            sys.stdout.write(f"{v}\n")
            sys.stdout.flush()
            self.quit(ExitCode.OK)


class ControlXpraClient(CommandConnectClient):
    """ Allows us to send commands to a server.
    """

    def set_command_args(self, command) -> None:
        self.command = command

    def timeout(self, *_args) -> None:
        self.warn_and_quit(ExitCode.TIMEOUT, "timeout: server did not respond")

    def do_command(self, caps: typedict) -> None:
        cr = caps.tupleget("command_response")
        if cr is None:
            self.warn_and_quit(ExitCode.UNSUPPORTED, "server does not support control command")
            return
        code = int(cr[0])
        text = str(cr[1])
        if code != 0:
            log.warn("server returned error code %s", code)
            self.warn_and_quit(ExitCode.REMOTE_ERROR, " %s" % text)
            return
        self.warn_and_quit(ExitCode.OK, text)

    def make_hello(self) -> dict[str, Any]:
        capabilities = super().make_hello()
        log("make_hello() adding command request '%s' to %s", self.command, capabilities)
        capabilities["command_request"] = tuple(self.command)
        capabilities["request"] = "command"
        return capabilities


class PrintClient(SendCommandConnectClient):
    """ Allows us to send a file to the server for printing.
    """

    def set_command_args(self, command) -> None:
        log("set_command_args(%s)", command)
        self.filename = command[0]
        # print command arguments:
        # filename, file_data, mimetype, source_uuid, title, printer, no_copies, print_options_str = packet[1:9]
        self.command = command[1:]

        # TODO: load as needed...

        if self.filename == "-":
            # replace with filename proposed
            self.filename = command[2]
            # read file from stdin
            with open(sys.stdin.fileno(), mode="rb", closefd=False) as stdin_binary:
                self.file_data = stdin_binary.read()
            log("read %i bytes from stdin", len(self.file_data))
        else:
            if not self.check_file_size(os.path.getsize(self.filename)):
                return
            from xpra.util.io import load_binary_file
            self.file_data = load_binary_file(self.filename)
            log("read %i bytes from %s", len(self.file_data), self.filename)
        if not self.check_file_size(len(self.file_data)):
            return
        assert self.file_data, "no data found for '%s'" % self.filename

    def check_file_size(self, size: int) -> bool:
        if size <= self.file_size_limit:
            return True
        self.warn_and_quit(ExitCode.FILE_TOO_BIG,
                           "the file is too large: %sB (the file size limit is %sB)" % (
                               std_unit(size), std_unit(self.file_size_limit)))
        return False

    def client_type(self) -> str:
        return "Python/GObject/Print"

    def timeout(self, *_args) -> None:
        self.warn_and_quit(ExitCode.TIMEOUT, "timeout: server did not respond")

    def do_command(self, caps: typedict) -> None:
        printing = caps.boolget("printing")
        if not printing:
            self.warn_and_quit(ExitCode.UNSUPPORTED, "server does not support printing")
            return
        # we don't compress file data
        # (this should run locally most of the time anyway)
        from xpra.net.compression import Compressed  # pylint: disable=import-outside-toplevel
        blob = Compressed("print", self.file_data)
        self.send("print", self.filename, blob, *self.command)
        log("print: sending %s as %s for printing", self.filename, blob)
        self.idle_add(self.send, "disconnect", ConnectionMessage.DONE.value, "detaching")

    def make_hello(self) -> dict[str, Any]:
        capabilities = super().make_hello()
        capabilities.setdefault("wants", []).append("features")
        capabilities["request"] = "print"
        return capabilities


class ExitXpraClient(HelloRequestClient):
    """ This client does one thing only:
        it asks the server to terminate (like stop),
        but without killing the Xvfb or clients.
    """

    def hello_request(self) -> dict[str, Any]:
        return {
            "request": "exit",
        }

    def do_command(self, caps: typedict) -> None:
        self.idle_add(self.send,
                      "exit-server",
                      os.environ.get("XPRA_EXIT_MESSAGE", ConnectionMessage.SERVER_EXIT.value)
                      )


class StopXpraClient(HelloRequestClient):
    """ stop a server """

    def hello_request(self) -> dict[str, Any]:
        return {
            "request": "stop",
        }

    def do_command(self, caps: typedict) -> None:
        if not self.server_client_shutdown:
            log.error("Error: cannot shutdown this server")
            log.error(" the feature is disabled on the server")
            self.quit(ExitCode.UNSUPPORTED)
            return
        # with newer (v5) servers, the "stop" request should have done it,
        # but let's try harder for older servers:
        self.timeout_add(1000, self.send_shutdown_server)
        # self.idle_add(self.send_shutdown_server)
        # not exiting the client here,
        # the server should send us the shutdown disconnection message anyway
        # and if not, we will then hit the timeout to tell us something went wrong


class DetachXpraClient(HelloRequestClient):
    """ run the `detach` subcommand """

    def hello_request(self) -> dict[str, Any]:
        return {
            "request": "detach",
        }

    def do_command(self, caps: typedict) -> None:
        self.idle_add(self.send, "disconnect", ConnectionMessage.DONE.value, "detaching")
        # not exiting the client here,
        # the server should disconnect us with the response


class WaitForDisconnectXpraClient(DetachXpraClient):
    """ we just want the connection to close """

    def _process_disconnect(self, _packet: PacketType) -> None:
        self.quit(ExitCode.OK)


class RequestStartClient(HelloRequestClient):
    """ request the system proxy server to start a new session for us """
    # wait longer for this command to return:
    from xpra.scripts.main import WAIT_SERVER_TIMEOUT
    COMMAND_TIMEOUT = EXTRA_TIMEOUT + WAIT_SERVER_TIMEOUT

    def dots(self) -> bool:
        errwrite(".")
        return not self.connection_established

    def _process_connection_lost(self, packet: PacketType) -> None:
        errwrite("\n")
        super()._process_connection_lost(packet)

    def hello_request(self) -> dict[str, Any]:
        if first_time("hello-request"):
            # this can be called again if we receive a challenge,
            # but only print this message once:
            errwrite("requesting new session, please wait")
        self.timeout_add(1 * 1000, self.dots)
        return {
            "start-new-session": self.start_new_session,
            # tells proxy servers we don't want to connect to the real / new instance:
            "connect": False,
        }

    def server_connection_established(self, caps: typedict) -> bool:
        # the server should respond with the display chosen
        log("server_connection_established() exit_code=%s", self.exit_code)
        display = caps.strget("display")
        if display:
            mode = caps.strget("mode")
            session_type = {
                "start": "seamless ",
                "start-desktop": "desktop ",
                "shadow": "shadow ",
            }.get(mode, "")
            try:
                errwrite("\n%ssession now available on display %s\n" % (session_type, display))
                if POSIX and not OSX and self.displayfd > 0 and display and display.startswith(":"):
                    from xpra.platform.displayfd import write_displayfd
                    log("writing display %s to displayfd=%s", display, self.displayfd)
                    write_displayfd(self.displayfd, display[1:])
            except OSError:
                log("server_connection_established(..)", exc_info=True)
        if not self.exit_code:
            self.quit(0)
        return True

    def __init__(self, opts):
        super().__init__(opts)
        try:
            self.displayfd = int(opts.displayfd)
        except (ValueError, TypeError):
            self.displayfd = 0
