#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Fenrir TTY screen reader
# By Chrys, Storm Dragon, and contributors.

import os
import os.path
import select
import socket
import struct
import time

from fenrirscreenreader.core import debug
from fenrirscreenreader.core import remoteInstanceRegistry
from fenrirscreenreader.core.eventData import FenrirEventType
from fenrirscreenreader.core.remoteDriver import RemoteDriver as remoteDriver


MAIN_SOCKET_FILE = "/tmp/fenrirscreenreader-daemon.sock"


class driver(remoteDriver):
    def __init__(self):
        remoteDriver.__init__(self)
        self.fenrirSocks = []
        self.bound_sockets = []

    def initialize(self, environment):
        self.env = environment
        # Use threading instead of multiprocessing to avoid pickle issues
        # with self.env (which contains unpicklable file handles)
        self.env["runtime"]["ProcessManager"].add_custom_event_thread(
            self.watch_dog, multiprocess=False
        )

    def _get_configured_socket_file(self):
        socket_file = ""
        try:
            socket_file = self.env["runtime"]["SettingsManager"].get_setting(
                "remote", "socket_file"
            )
        except Exception as e:
            self.env["runtime"]["DebugManager"].write_debug_out(
                "unixDriver watch_dog: Error getting socket file setting: "
                + str(e),
                debug.DebugLevel.ERROR,
            )
        return socket_file

    def _get_socket_candidates(self):
        configured_socket_file = self._get_configured_socket_file()
        if configured_socket_file:
            return [(configured_socket_file, False)]

        screen_driver = self.env["runtime"]["SettingsManager"].get_setting(
            "screen", "driver"
        )
        if screen_driver == "vcsaDriver":
            return [(MAIN_SOCKET_FILE, False)]

        private_socket_file = (
            "/tmp/fenrirscreenreader-" + str(os.getpid()) + ".sock"
        )
        return [(private_socket_file, False), (MAIN_SOCKET_FILE, True)]

    def _is_socket_active(self, socket_file):
        test_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        try:
            test_sock.settimeout(0.2)
            test_sock.connect(socket_file)
            return True
        except OSError:
            return False
        finally:
            test_sock.close()

    def _bind_socket(self, socket_file, optional):
        if os.path.exists(socket_file):
            if optional and self._is_socket_active(socket_file):
                return None
            os.unlink(socket_file)

        fenrir_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        try:
            fenrir_sock.bind(socket_file)
            os.chmod(socket_file, 0o666)  # Allow all users to read/write
            fenrir_sock.listen(1)
        except OSError:
            fenrir_sock.close()
            if optional:
                return None
            raise
        return fenrir_sock

    def _register_instance(self):
        settings_manager = self.env["runtime"]["SettingsManager"]
        active_socket_files = [
            socket_file
            for fenrir_sock, socket_file in self.bound_sockets
            if self._socket_file_matches_socket(fenrir_sock, socket_file)
        ]
        instance_data = {
            "pid": os.getpid(),
            "ppid": os.getppid(),
            "socket_files": active_socket_files,
            "main_socket": MAIN_SOCKET_FILE in active_socket_files,
            "screen_driver": settings_manager.get_setting("screen", "driver"),
            "keyboard_driver": settings_manager.get_setting("keyboard", "driver"),
            "x11_window_id": settings_manager.get_setting(
                "keyboard", "x11_window_id"
            ),
            "created_at": time.time(),
        }
        remoteInstanceRegistry.write_instance(instance_data)

    def _socket_file_matches_socket(self, fenrir_sock, socket_file):
        try:
            if os.stat(socket_file).st_uid != os.getuid():
                return False
            fd_stat = os.fstat(fenrir_sock.fileno())
            socket_inode = str(fd_stat.st_ino)
            with open("/proc/net/unix", "r", encoding="utf-8") as proc_file:
                for line in proc_file:
                    line_parts = line.split()
                    if (
                        len(line_parts) >= 8
                        and line_parts[6] == socket_inode
                        and line_parts[7] == socket_file
                    ):
                        return True
        except OSError:
            return False
        return False

    def _socket_file_for_socket(self, fenrir_sock):
        for bound_sock, socket_file in self.bound_sockets:
            if bound_sock == fenrir_sock:
                return socket_file
        return ""

    def _get_peer_pid(self, client_sock):
        so_peercred = getattr(socket, "SO_PEERCRED", 17)
        try:
            creds = client_sock.getsockopt(
                socket.SOL_SOCKET, so_peercred, struct.calcsize("3i")
            )
            pid, _uid, _gid = struct.unpack("3i", creds)
            return pid
        except OSError:
            return 0

    def _get_parent_pid(self, pid):
        try:
            with open(f"/proc/{pid}/stat", "r", encoding="utf-8") as proc_file:
                stat_text = proc_file.read()
        except OSError:
            return 0

        end_command = stat_text.rfind(")")
        if end_command == -1:
            return 0
        stat_fields = stat_text[end_command + 2 :].split()
        if len(stat_fields) < 2:
            return 0
        try:
            return int(stat_fields[1])
        except ValueError:
            return 0

    def _find_ancestor_private_socket(self, pid):
        seen_pids = set()
        while pid > 1 and pid not in seen_pids:
            seen_pids.add(pid)
            socket_file = f"/tmp/fenrirscreenreader-{pid}.sock"
            if self._is_registered_private_socket(socket_file):
                return socket_file
            pid = self._get_parent_pid(pid)
        return ""

    def _is_registered_private_socket(self, socket_file):
        for instance in remoteInstanceRegistry.list_instances():
            if socket_file in instance.get("socket_files", []):
                return True
        return False

    def _get_registered_private_sockets(self):
        socket_files = []
        for instance in remoteInstanceRegistry.list_instances():
            for socket_file in instance.get("socket_files", []):
                if socket_file == MAIN_SOCKET_FILE:
                    continue
                if socket_file in socket_files:
                    continue
                socket_files.append(socket_file)
        return socket_files

    def _find_available_private_socket(self, preferred_socket=""):
        socket_files = self._get_registered_private_sockets()
        if preferred_socket and preferred_socket in socket_files:
            socket_files.remove(preferred_socket)
            socket_files.insert(0, preferred_socket)

        for socket_file in socket_files:
            if self._is_own_socket_file(socket_file):
                return socket_file
            if self._is_socket_active(socket_file):
                return socket_file
        return ""

    def _is_own_socket_file(self, socket_file):
        return any(
            socket_file == bound_socket_file
            for _bound_sock, bound_socket_file in self.bound_sockets
        )

    def _has_own_private_socket(self):
        return any(
            bound_socket_file != MAIN_SOCKET_FILE
            for _bound_sock, bound_socket_file in self.bound_sockets
        )

    def _forward_remote_to_socket(self, data, socket_file):
        forward_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        try:
            forward_sock.settimeout(0.2)
            forward_sock.connect(socket_file)
            forward_sock.sendall((data + "\n").encode("utf-8"))
            return True
        except OSError as e:
            self.env["runtime"]["DebugManager"].write_debug_out(
                "unixDriver watch_dog: Error forwarding remote data to "
                + socket_file
                + ": "
                + str(e),
                debug.DebugLevel.ERROR,
            )
            return False
        finally:
            forward_sock.close()

    def _route_main_socket_command(self, data, client_sock, socket_file):
        if socket_file != MAIN_SOCKET_FILE:
            return False
        if not self._has_own_private_socket():
            return False

        peer_pid = self._get_peer_pid(client_sock)
        ancestor_socket = ""
        if peer_pid > 1:
            ancestor_socket = self._find_ancestor_private_socket(peer_pid)
        target_socket = self._find_available_private_socket(ancestor_socket)
        if not target_socket:
            return False

        if self._is_own_socket_file(target_socket):
            return False

        return self._forward_remote_to_socket(data, target_socket)

    def _cleanup(self):
        for fenrir_sock in self.fenrirSocks:
            try:
                fenrir_sock.close()
            except OSError:
                pass
        self.fenrirSocks = []

        for _, socket_file in self.bound_sockets:
            try:
                if (
                    os.path.exists(socket_file)
                    and not self._is_socket_active(socket_file)
                ):
                    os.unlink(socket_file)
            except OSError:
                pass
        self.bound_sockets = []
        remoteInstanceRegistry.remove_instance()

    def _handle_client(self, client_sock, event_queue, socket_file=""):
        try:
            rawdata = client_sock.recv(8129)
        except Exception as e:
            self.env["runtime"]["DebugManager"].write_debug_out(
                "unixDriver watch_dog: Error receiving data from client: "
                + str(e),
                debug.DebugLevel.ERROR,
            )
            rawdata = b""

        try:
            data = rawdata.decode("utf-8").rstrip().lstrip()
            upper_data = data.upper()
            if upper_data in ("LS", "LIST", "COMMAND LS", "COMMAND LIST"):
                response = self.env["runtime"][
                    "RemoteManager"
                ].handle_remote_incomming_with_response(data)
                client_sock.sendall((response["message"] + "\n").encode("utf-8"))
                return

            if self._route_main_socket_command(data, client_sock, socket_file):
                return

            event_queue.put(
                {
                    "Type": FenrirEventType.remote_incomming,
                    "data": data,
                }
            )
        except Exception as e:
            self.env["runtime"]["DebugManager"].write_debug_out(
                "unixDriver watch_dog: Error decoding/queuing data: "
                + str(e),
                debug.DebugLevel.ERROR,
            )

    def watch_dog(self, active, event_queue):
        # echo "command say this is a test" | socat -
        # UNIX-CLIENT:/tmp/fenrirscreenreader-daemon.sock
        for socket_file, optional in self._get_socket_candidates():
            fenrir_sock = self._bind_socket(socket_file, optional)
            if fenrir_sock is None:
                continue
            self.fenrirSocks.append(fenrir_sock)
            self.bound_sockets.append((fenrir_sock, socket_file))

        if not self.fenrirSocks:
            return

        self._register_instance()
        last_register = time.time()
        while active.value:
            if time.time() - last_register > 10.0:
                self._register_instance()
                last_register = time.time()

            # Check if the client is still connected and if data is available:
            try:
                r, _, _ = select.select(self.fenrirSocks, [], [], 0.8)
            except select.error:
                break
            if r == []:
                continue
            for fenrir_sock in r:
                client_sock, client_addr = fenrir_sock.accept()
                socket_file = self._socket_file_for_socket(fenrir_sock)
                # Ensure client socket is always closed to prevent resource
                # leaks
                try:
                    self._handle_client(client_sock, event_queue, socket_file)
                finally:
                    # Always close client socket, even if data processing fails
                    try:
                        client_sock.close()
                    except Exception as e:
                        self.env["runtime"]["DebugManager"].write_debug_out(
                            "unixDriver watch_dog: Error closing client socket: "
                            + str(e),
                            debug.DebugLevel.ERROR,
                        )
        self._cleanup()

    def shutdown(self):
        self._cleanup()
        remoteDriver.shutdown(self)
