import contextlib
import ctypes
import dataclasses
import logging
import os
import sys
import threading
import typing
from collections import namedtuple
from enum import Enum, auto
from functools import cached_property, wraps
from pathlib import Path
from select import select
from typing import Any, Optional

from xonsh.built_ins import XSH
from xonsh.main import main as xonsh_main

from rpcclient.clients.darwin.consts import BLOCK_IS_GLOBAL
from rpcclient.core.capture_fd import CaptureFD
from rpcclient.core.structs.consts import (
    EAGAIN,
    ECONNREFUSED,
    EEXIST,
    EISDIR,
    ENOENT,
    ENOTDIR,
    ENOTEMPTY,
    EPERM,
    EPIPE,
    RTLD_NEXT,
)
from rpcclient.core.structs.generic import block_descriptor, block_literal
from rpcclient.core.subsystems.decorator import subsystem
from rpcclient.core.subsystems.fs import Fs
from rpcclient.core.subsystems.lief import Lief
from rpcclient.core.subsystems.network import Network
from rpcclient.core.subsystems.processes import Processes
from rpcclient.core.subsystems.sysctl import Sysctl
from rpcclient.core.symbol import Symbol
from rpcclient.core.symbols_jar import SymbolsJar
from rpcclient.event_notifier import EventNotifier
from rpcclient.exceptions import (
    ArgumentError,
    BadReturnValueError,
    RpcBrokenPipeError,
    RpcConnectionRefusedError,
    RpcFileExistsError,
    RpcFileNotFoundError,
    RpcIsADirectoryError,
    RpcNotADirectoryError,
    RpcNotEmptyError,
    RpcPermissionError,
    RpcResourceTemporarilyUnavailableError,
    ServerResponseError,
    SpawnError,
)
from rpcclient.protocol.rpc_bridge import RpcBridge
from rpcclient.protos.rpc_api_pb2 import Argument, MsgId
from rpcclient.protos.rpc_pb2 import ProtocolConstants

tty_support = False
try:
    import termios
    import tty

    tty_support = True
except ImportError:
    logging.warning("termios not available on your system. some functionality may not work as expected")

io_or_str = typing.TypeVar("io_or_str", typing.IO, str)

SpawnResult = namedtuple("SpawnResult", "error pid stdout")

INVALID_PID = 0xFFFFFFFF
CHUNK_SIZE = 1024

USAGE = """
Welcome to the rpcclient interactive shell! You interactive shell for controlling the remote rpcserver.
Feel free to use the following globals:

🌍 p - the injected process
🌍 symbols - process global symbols

Have a nice flight ✈️!
Starting an IPython shell... 🐍
"""


@dataclasses.dataclass
class ProtocolDitentStat:
    errno: int
    st_dev: int  # device inode resides on
    st_mode: int  # inode protection mode
    st_nlink: int  # number of hard links to the file
    st_ino: int  # inode's number
    st_uid: int  # user-id of owner
    st_gid: int  # group-id of owner
    st_rdev: int  # device type, for special file inode
    st_size: int  # file size, in bytes
    st_blocks: int  # blocks allocated for file
    st_blksize: int  # optimal blocksize for I/O
    st_atime: int
    st_mtime: int
    st_ctime: int


@dataclasses.dataclass
class ProtocolDirent:
    d_inode: int
    d_type: int
    d_name: str
    lstat: ProtocolDitentStat
    stat: ProtocolDitentStat


class ClientEvent(Enum):
    CREATED = auto()
    TERMINATED = auto()


def null_pointer_guard(func: typing.Callable) -> typing.Callable:
    """
    A decorator to prevent dereferencing a null pointer by checking if the given
    address is zero before calling the wrapped function. If the address is zero,
    an ArgumentError is raised.

    Parameters:
        func (Callable): The function to be wrapped by the decorator.

    Returns:
        Callable: The wrapped function with null pointer guard applied.

    Raises:
        ArgumentError: If the 'address' argument is zero.
    """

    @wraps(func)
    def wrapper(self: typing.Any, address: int, *args: typing.Any, **kwargs: typing.Any) -> typing.Any:
        if address == 0:
            raise ArgumentError("Unable dereference null pointer.")
        return func(self, address, *args, **kwargs)

    return wrapper


class CoreClient:
    """Main client interface to access the remote rpcserver"""

    DEFAULT_ARGV: typing.ClassVar = ["/bin/sh"]
    DEFAULT_ENVP: typing.ClassVar = []

    def __init__(self, bridge: RpcBridge, dlsym_global_handle=RTLD_NEXT):
        self._bridge = bridge
        self._old_settings = None
        self._endianness = "<"
        self._dlsym_global_handle = dlsym_global_handle
        self._protocol_lock = threading.Lock()
        self._logger = logging.getLogger(self.__module__)
        self.notifier = EventNotifier()
        self.symbols = SymbolsJar.create(self)

    @subsystem
    def fs(self) -> Fs:
        return Fs(self)

    @subsystem
    def processes(self) -> Processes:
        return Processes(self)

    @subsystem
    def network(self) -> Network:
        return Network(self)

    @subsystem
    def lief(self) -> Lief:
        return Lief(self)

    @subsystem
    def sysctl(self) -> Sysctl:
        return Sysctl(self)

    def info(self):
        """print information about the current target"""
        uname = self.uname
        print("sysname:", uname.sysname)
        print("nodename:", uname.nodename)
        print("release:", uname.release)
        print("version:", uname.version)
        print("machine:", uname.machine)
        print(f"uid: {self.symbols.getuid():d}")
        print(f"gid: {self.symbols.getgid():d}")
        print(f"pid: {self.pid:d}")
        print(f"ppid: {self.symbols.getppid():d}")
        print(f"progname: {self.progname}")

    @cached_property
    def progname(self) -> str:
        """get the program name from remote"""
        # bugfix https://github.com/doronz88/rpc-project/issues/405
        # Default Linux libraries don't expose `getprogname()`, so instead we use the `__progname` global symbol.
        # Tested on both macOS and Ubuntu.
        return self.symbols["__progname"][0].peek_str()

    @cached_property
    def uname(self):
        """get the utsname struct from remote"""
        raise NotImplementedError()

    @property
    def id(self) -> int:
        return self._bridge.client_id

    @property
    def platform(self) -> str:
        return self._bridge.platform

    @property
    def sysname(self) -> str:
        return self._bridge.sysname

    @property
    def arch(self) -> int:
        return self._bridge.arch

    def rpc_call(self, msg_id: int, **kwargs):
        try:
            return self._bridge.rpc_call(msg_id, **kwargs)
        except ConnectionError:
            self.notifier.notify(ClientEvent.TERMINATED, self.id)
            raise
        except ServerResponseError:
            raise

    def dlopen(self, filename: str, mode: int) -> Symbol:
        """Load a shared library on the remote host and return its handle."""
        return self.symbol(self.rpc_call(MsgId.REQ_DLOPEN, filename=filename, mode=mode).handle)

    def dlclose(self, lib: int) -> int:
        """Close a previously opened remote library handle."""
        return self.rpc_call(MsgId.REQ_DLCLOSE, handle=ctypes.c_uint64(lib).value).res

    def dlsym(self, lib: int, symbol_name: str) -> int:
        """Resolve a symbol name in a remote library handle and return its address."""
        return self.rpc_call(MsgId.REQ_DLSYM, handle=ctypes.c_uint64(lib).value, symbol_name=symbol_name).ptr

    @null_pointer_guard
    def call(
        self,
        address: int,
        argv: Optional[list[int]] = None,
        return_float64=False,
        return_float32=False,
        return_raw=False,
        va_list_index: int = 0xFFFF,
    ) -> typing.Union[float, Symbol, Any]:
        """call a remote function and retrieve its return value as a Symbol object"""
        args = []
        for arg in argv:
            if isinstance(arg, float):
                args.append(Argument(v_double=arg))
            elif isinstance(arg, str):
                args.append(Argument(v_str=arg))
            elif isinstance(arg, int):
                args.append(Argument(v_int=ctypes.c_uint64(arg).value))
            elif isinstance(arg, bytes):
                args.append(Argument(v_bytes=arg))
            elif isinstance(arg, Enum):
                args.append(Argument(v_int=ctypes.c_uint64(arg.value).value))
            else:
                raise ArgumentError()

        ret = self.rpc_call(MsgId.REQ_CALL, address=address, va_list_index=va_list_index, argv=args)
        if ret.HasField("arm_registers"):
            d0 = ret.arm_registers.d0
            if return_float32:
                return ctypes.c_float(d0).value
            if return_float64:
                return d0
            if return_raw:
                return ret.arm_registers
            return self.symbol(ret.arm_registers.x0)
        return self.symbol(ret.return_value)

    @null_pointer_guard
    def peek(self, address: int, size: int) -> bytes:
        """peek data at the given address"""
        try:
            return self.rpc_call(MsgId.REQ_PEEK, address=address, size=size).data
        except ServerResponseError as e:
            raise ArgumentError() from e

    @null_pointer_guard
    def poke(self, address: int, data: bytes):
        """poke data at a given address"""
        try:
            return self.rpc_call(MsgId.REQ_POKE, address=address, data=data)
        except ServerResponseError as e:
            raise ArgumentError() from e

    def get_dummy_block(self) -> Symbol:
        """Get an address for a stub block containing nothing"""
        block_size = block_literal.sizeof()
        desc_size = block_descriptor.sizeof()

        descriptor = self.symbols.malloc(desc_size)
        descriptor.poke(block_descriptor.build({"reserved": 0, "size": block_size}))

        block = self.symbols.malloc(block_size)
        block.poke(
            block_literal.build({
                "isa": self.symbols._NSConcreteGlobalBlock,
                "flags": BLOCK_IS_GLOBAL,
                "reserved": 0,
                "invoke": self.symbols.getpid,
                "descriptor": descriptor,
            })
        )
        return block

    def listdir(self, path: str):
        """get an address for a stub block containing nothing"""
        entries = []
        try:
            ret = self.rpc_call(MsgId.REQ_LIST_DIR, path=path)
        except ServerResponseError:
            self.raise_errno_exception(f"failed to listdir: {path}")

        for entry in ret.dir_entries:
            lstat = ProtocolDitentStat(
                errno=entry.lstat.errno1,
                st_blocks=entry.lstat.st_blocks,
                st_blksize=entry.lstat.st_blksize,
                st_atime=entry.lstat.st_atime1,
                st_ctime=entry.lstat.st_ctime1,
                st_mtime=entry.lstat.st_mtime1,
                st_nlink=entry.lstat.st_nlink,
                st_mode=entry.lstat.st_mode,
                st_rdev=entry.lstat.st_rdev,
                st_size=entry.lstat.st_size,
                st_dev=entry.lstat.st_dev,
                st_gid=entry.lstat.st_gid,
                st_ino=entry.lstat.st_ino,
                st_uid=entry.lstat.st_uid,
            )
            stat = ProtocolDitentStat(
                errno=entry.stat.errno1,
                st_blocks=entry.stat.st_blocks,
                st_blksize=entry.stat.st_blksize,
                st_atime=entry.stat.st_atime1,
                st_ctime=entry.stat.st_ctime1,
                st_mtime=entry.stat.st_mtime1,
                st_nlink=entry.stat.st_nlink,
                st_mode=entry.stat.st_mode,
                st_rdev=entry.stat.st_rdev,
                st_size=entry.stat.st_size,
                st_dev=entry.stat.st_dev,
                st_gid=entry.stat.st_gid,
                st_ino=entry.stat.st_ino,
                st_uid=entry.stat.st_uid,
            )
            entries.append(
                ProtocolDirent(
                    d_inode=entry.lstat.st_ino, d_type=entry.d_type, d_name=entry.d_name, lstat=lstat, stat=stat
                )
            )
        return entries

    def spawn(
        self,
        argv: Optional[list[str]] = None,
        envp: Optional[list[str]] = None,
        stdin: io_or_str = sys.stdin,
        stdout=sys.stdout,
        raw_tty=False,
        background=False,
    ) -> SpawnResult:
        """
        spawn a new process and forward its stdin, stdout & stderr

        :param argv: argv of the process to be executed
        :param envp: envp of the process to be executed
        :param stdin: either a file object to read from OR a string
        :param stdout: a file object to write both stdout and stderr to. None if background is requested
        :param raw_tty: should enable raw tty mode
        :param background: should execute process in background
        :return: a SpawnResult. error is None if background is requested
        """
        if argv is None:
            argv = self.DEFAULT_ARGV

        if envp is None:
            envp = self.DEFAULT_ENVP

        with self._protocol_lock:
            pid = self._execute(argv, envp, background=background)
            self._logger.info(f"shell process started as pid: {pid}")

            if background:
                return SpawnResult(error=None, pid=pid, stdout=None)

            if raw_tty:
                self._prepare_terminal()
            try:
                error = self.enter_pty_mode(stdin, stdout)
            except Exception:
                # this is important to really catch every exception here, even exceptions not inheriting from Exception
                # so the controlling terminal will remain working with its previous settings
                if raw_tty:
                    self._restore_terminal()
                raise

        return SpawnResult(error=error, pid=pid, stdout=stdout)

    def symbol(self, symbol: int):
        """at a symbol object from a given address"""
        return Symbol.create(symbol, self)

    @property
    def errno(self) -> int:
        return self.symbols.errno[0]

    @errno.setter
    def errno(self, value: int) -> None:
        self.symbols.errno[0] = value

    @property
    def last_error(self):
        """get info about the last occurred error"""
        if not self.errno:
            return ""
        err_str = self.symbols.strerror(self.errno).peek_str()
        return f"[{self.errno}] {err_str}"

    @property
    def environ(self) -> list[str]:
        result = []
        environ = self.symbols.environ[0]
        i = 0
        while environ[i]:
            result.append(environ[i].peek_str())
            i += 1
        return result

    def setenv(self, name: str, value: str):
        """set process environment variable"""
        self.symbols.setenv(name, value)

    def getenv(self, name: str) -> typing.Optional[str]:
        """get process environment variable"""
        value = self.symbols.getenv(name)
        if not value:
            return None
        return value.peek_str()

    @cached_property
    def pid(self):
        return int(self.symbols.getpid())

    @contextlib.contextmanager
    def safe_calloc(self, size: int):
        with self.safe_malloc(size) as x:
            x.poke(b"\x00" * size)
            try:
                yield x
            finally:
                pass

    @contextlib.contextmanager
    def safe_malloc(self, size: int):
        with self.freeing(self.symbols.malloc(size)) as x:
            yield x

    @contextlib.contextmanager
    def freeing(self, symbol):
        try:
            yield symbol
        finally:
            if symbol:
                self.symbols.free(symbol)

    def close(self):
        self.rpc_call(MsgId.REQ_CLOSE_CLIENT)
        self.notifier.notify(ClientEvent.TERMINATED, self.id)
        self._bridge.close()

    def shell(self):
        self._logger.disabled = True

        args = ["--rc"]
        args.append(str((Path(__file__).parent / "xonshrc.py").absolute()))

        XSH.ctx["_client_to_reuse"] = self

        try:
            logging.getLogger("parso.python.diff").disabled = True
            logging.getLogger("parso.cache").disabled = True
            logging.getLogger("asyncio").disabled = True
            xonsh_main(args)
        except SystemExit:
            self._logger.disabled = False

    def _execute(self, argv: list[str], envp: list[str], background=False) -> int:
        try:
            return self.rpc_call(MsgId.REQ_EXEC, background=background, argv=argv, envp=envp).pid
        except ServerResponseError as e:
            raise SpawnError(f"failed to spawn: {argv}") from e

    def enter_pty_mode(self, stdin=sys.stdin, stdout=sys.stdout):
        # the socket must be non-blocking for using select()
        sock = self._bridge.sock.raw_socket
        sock.setblocking(False)
        exit_code = None
        try:
            fds = []
            if hasattr(stdin, "fileno"):
                fds.append(stdin)
            else:
                data = stdin
                if isinstance(data, str):
                    data = data.encode()
                sock.sendall(data)
            fds.append(sock)

            running = True
            while running:
                rlist, _, _ = select(fds, [], [])
                for fd in rlist:
                    if fd == stdin or (stdin is sys.stdin and fd == sys.stdin):
                        if stdin is sys.stdin:
                            buf = os.read(stdin.fileno(), ProtocolConstants.RPC_PTY_BUFFER_SIZE)
                        else:
                            buf = stdin.read(ProtocolConstants.RPC_PTY_BUFFER_SIZE)
                            if isinstance(buf, str):
                                buf = buf.encode()
                        if buf:
                            sock.sendall(buf)
                    elif fd == sock:
                        try:
                            response = self._bridge.sock.rpc_msg_recv_pty()
                        except ConnectionResetError:
                            print("Bye. 👋")
                            running = False
                            break
                        msg_type = response.WhichOneof("type")
                        if msg_type == "buffer":
                            stdout.write(response.buffer.decode())
                            if hasattr(stdout, "flush"):
                                stdout.flush()
                        elif msg_type == "exit_code":
                            exit_code = response.exit_code
                            running = False
                            break
        finally:
            sock.setblocking(True)
        return exit_code

    def _restore_terminal(self):
        if not tty_support:
            return
        termios.tcsetattr(sys.stdin.fileno(), termios.TCSADRAIN, self._old_settings)

    def _prepare_terminal(self):
        if not tty_support:
            return
        fd = sys.stdin
        self._old_settings = termios.tcgetattr(fd)
        tty.setraw(fd)

    def raise_errno_exception(self, message: str):
        message += f" ({self.last_error})"
        exceptions = {
            EPERM: RpcPermissionError,
            ENOENT: RpcFileNotFoundError,
            EEXIST: RpcFileExistsError,
            EISDIR: RpcIsADirectoryError,
            ENOTDIR: RpcNotADirectoryError,
            EPIPE: RpcBrokenPipeError,
            ENOTEMPTY: RpcNotEmptyError,
            EAGAIN: RpcResourceTemporarilyUnavailableError,
            ECONNREFUSED: RpcConnectionRefusedError,
        }
        exception = exceptions.get(self.errno)
        if exception:
            raise exception(message)
        raise BadReturnValueError(message)

    def capture_fd(self, fd: int, sock_buf_size: Optional[int] = None) -> CaptureFD:
        """
        Get a context manager, capturing output to `fd`. Read from it using the `read()` method

        sock_buf_size is required for captures above 6KB, as any write above this value would block until a read is performed.

        :param fd: FD to capture
        :param sock_buf_size: Buffer size for the capture socket, if not specified, default value is used.
        :return: CaptureFD object
        """
        return CaptureFD(self, fd, sock_buf_size)

    def __repr__(self):
        return f"<{self.__class__.__name__}: {self.pid} | {self.progname}>"
