import asyncio
import os
import shlex
import shutil
import signal
import sys

import asyncssh

from ...base import Result
from ...log import DEBUG, ERR, NOTICE, log
from ..SSHClient import SSHClient as Base
from .util import join_cmd

_USE_DEFAULT_KNOWN_HOSTS = object()

class AsyncSSH(Base):

    def __init__(
        self,
        uri: str,
        *,
        client_keys: list[str] | None = None,
        known_hosts = _USE_DEFAULT_KNOWN_HOSTS,
        term_type: str | None = None,
        connect_timeout: float | None = 30.0,
        **kwargs,
    ) -> None:

        super().__init__(
            uri,
            caps = self.Caps.LogOutput
            | self.Caps.Wd
            | self.Caps.Interactive
            | self.Caps.ModEnv,
            **kwargs,
        )

        self.__client_keys = client_keys
        self.__known_hosts = known_hosts
        self.__term_type = term_type or os.environ.get('TERM', 'xterm')
        self.__connect_timeout = connect_timeout
        self.__conn: asyncssh.SSHClientConnection | None = None

    async def _open(self) -> None:
        await super()._open()
        await self._conn

    async def _close(self) -> None:
        if self.__conn is not None:
            try:
                self.__conn.close()
                await self.__conn.wait_closed()
            except Exception as e:
                log(DEBUG, f'Failed to close connection ({str(e)}, ignored)')
            self.__conn = None

    def _connect_kwargs(self, hide_secrets: bool = False) -> dict:
        kwargs: dict = {
            'host': self.hostname,
            'port': self.port,
            'username': self.username,
            'password': self.password,
            'client_keys': self.__client_keys,
            'connect_timeout': self.__connect_timeout,
        }
        if self.__known_hosts is not _USE_DEFAULT_KNOWN_HOSTS:
            kwargs['known_hosts'] = self.__known_hosts
        ret = {k: v for k, v in kwargs.items() if v is not None}
        if hide_secrets and 'password' in kwargs:
            kwargs['password'] = '<hidden>'
        return ret

    @property
    async def _conn(self) -> asyncssh.SSHClientConnection:
        if self.__conn is None:
            try:
                self.__conn = await asyncssh.connect(**self._connect_kwargs())
            except Exception as e:
                msg = f'-------------------- Failed to connect ({str(e)})'
                log(ERR, ',', msg)
                for key, val in self._connect_kwargs(hide_secrets = True).items():
                    log(ERR, f'| {key:<20} = {val}')
                log(ERR, '`', msg)
                raise
        return self.__conn

    @staticmethod
    def _build_remote_command(cmd: list[str], wd: str | None) -> str:
        inner = f'exec {join_cmd(cmd)}'
        if wd is not None:
            inner = f'cd {shlex.quote(wd)} && {inner}'
        return f'/bin/sh -lc {shlex.quote(inner)}'

    @staticmethod
    def _has_local_tty() -> bool:
        try:
            return sys.stdin.isatty() and sys.stdout.isatty()
        except Exception:
            return False

    @staticmethod
    def _get_local_term_size() -> tuple[int, int, int, int]:
        cols, rows = shutil.get_terminal_size(fallback = (80, 24))
        xpixel = ypixel = 0
        try:
            import fcntl
            import struct
            import termios

            packed = fcntl.ioctl(sys.stdout.fileno(), termios.TIOCGWINSZ, b'\0' * 8)
            rows2, cols2, xpixel, ypixel = struct.unpack('HHHH', packed)
            if cols2 > 0 and rows2 > 0:
                cols, rows = cols2, rows2
        except Exception:
            pass
        return (cols, rows, xpixel, ypixel)

    async def _read_stream(
        self,
        stream,
        prio,
        collector: list[bytes],
        *,
        verbose: bool,
        log_prefix: str,
        log_enc: str,
    ) -> None:
        buf = b''
        while True:
            chunk = await stream.read(4096)
            if not chunk:
                break
            collector.append(chunk)
            if verbose:
                buf += chunk
                while b'\n' in buf:
                    line, buf = buf.split(b'\n', 1)
                    log(prio, log_prefix, line.decode(log_enc, errors = 'replace'))
        if verbose and buf:
            log(prio, log_prefix, buf.decode(log_enc, errors = 'replace'))

    async def _run_interactive_on_conn(
        self,
        cmd: list[str],
        wd: str | None,
        cmd_input: bytes | None,
        mod_env: dict[str, str] | None,
    ) -> Result:

        conn = await self._conn
        command = self._build_remote_command(cmd, wd)
        stdout_parts: list[bytes] = []

        proc = await conn.create_process(
            command = command,
            env = mod_env,
            stdin = asyncssh.PIPE,
            stdout = asyncssh.PIPE,
            stderr = asyncssh.STDOUT,
            encoding = None,
            request_pty = 'force',
            term_type = self.__term_type,
            term_size = self._get_local_term_size(),
        )

        loop = asyncio.get_running_loop()
        stdin_fd = sys.stdin.fileno()

        stdin_queue: asyncio.Queue[bytes | None] = asyncio.Queue()
        old_tty_state = None
        old_winch_handler = None
        stdin_reader_installed = False

        def _write_local(data: bytes) -> None:
            try:
                sys.stdout.buffer.write(data)
                sys.stdout.buffer.flush()
            except AttributeError:
                os.write(sys.stdout.fileno(), data)

        def _on_stdin_ready() -> None:
            try:
                data = os.read(stdin_fd, 4096)
            except OSError:
                data = b''

            if data:
                stdin_queue.put_nowait(data)
            else:
                try:
                    loop.remove_reader(stdin_fd)
                except Exception:
                    pass
                stdin_queue.put_nowait(None)

        async def _pump_stdin() -> None:
            if cmd_input is not None and proc.stdin is not None:
                proc.stdin.write(cmd_input)
                await proc.stdin.drain()
            while True:
                data = await stdin_queue.get()
                if data is None:
                    if proc.stdin is not None:
                        try:
                            proc.stdin.write_eof()
                        except (BrokenPipeError, OSError):
                            pass
                    return
                if proc.stdin is None:
                    return
                proc.stdin.write(data)
                await proc.stdin.drain()

        async def _pump_stdout() -> None:
            while True:
                chunk = await proc.stdout.read(4096)
                if not chunk:
                    break
                stdout_parts.append(chunk)
                _write_local(chunk)

        def _on_winch(*_args) -> None:

            try:
                proc.change_terminal_size(*self._get_local_term_size())
            except Exception:
                pass

        try:
            sys.stdout.flush()
            sys.stderr.flush()

            try:
                import termios
                import tty

                old_tty_state = termios.tcgetattr(stdin_fd)
                tty.setraw(stdin_fd)
            except Exception:
                old_tty_state = None

            try:
                loop.add_reader(stdin_fd, _on_stdin_ready)
                stdin_reader_installed = True
            except (NotImplementedError, RuntimeError):
                stdin_queue.put_nowait(None)

            if hasattr(signal, 'SIGWINCH'):
                try:
                    old_winch_handler = signal.getsignal(signal.SIGWINCH)
                    signal.signal(signal.SIGWINCH, _on_winch)
                except Exception:
                    old_winch_handler = None

            stdin_task = asyncio.create_task(_pump_stdin())
            stdout_task = asyncio.create_task(_pump_stdout())

            completed = await proc.wait(check = False)
            await stdout_task

            if not stdin_task.done():
                stdin_task.cancel()
                try:
                    await stdin_task
                except asyncio.CancelledError:
                    pass

            exit_code = completed.exit_status
            if exit_code is None:
                exit_code = (
                    completed.returncode if completed.returncode is not None else -1
                )

            stdout = b''.join(stdout_parts) if stdout_parts else None
            return Result(stdout, None, exit_code)

        finally:
            if stdin_reader_installed:
                try:
                    loop.remove_reader(stdin_fd)
                except Exception:
                    pass

            if old_winch_handler is not None and hasattr(signal, 'SIGWINCH'):
                try:
                    signal.signal(signal.SIGWINCH, old_winch_handler)
                except Exception:
                    pass

            if old_tty_state is not None:
                try:
                    import termios

                    termios.tcsetattr(stdin_fd, termios.TCSADRAIN, old_tty_state)
                except Exception:
                    pass

            try:
                sys.stdout.flush()
                sys.stderr.flush()
            except Exception:
                pass

    async def _run_captured_pty_on_conn(
        self,
        cmd: list[str],
        wd: str | None,
        verbose: bool,
        cmd_input: bytes | None,
        mod_env: dict[str, str] | None,
        log_prefix: str,
    ) -> Result:

        conn = await self._conn
        command = self._build_remote_command(cmd, wd)

        stdout_parts: list[bytes] = []
        stdout_log_enc = sys.stdout.encoding or 'utf-8'

        proc = await conn.create_process(
            command = command,
            env = mod_env,
            stdin = asyncssh.PIPE if cmd_input is not None else asyncssh.DEVNULL,
            stdout = asyncssh.PIPE,
            stderr = asyncssh.STDOUT,
            encoding = None,
            request_pty = 'force',
            term_type = self.__term_type,
        )

        task = asyncio.create_task(
            self._read_stream(
                proc.stdout,
                NOTICE,
                stdout_parts,
                verbose = verbose,
                log_prefix = log_prefix,
                log_enc = stdout_log_enc,
            )
        )

        if cmd_input is not None and proc.stdin is not None:
            proc.stdin.write(cmd_input)
            await proc.stdin.drain()
            proc.stdin.write_eof()

        completed = await proc.wait(check = False)
        await task

        exit_code = completed.exit_status
        if exit_code is None:
            exit_code = completed.returncode if completed.returncode is not None else -1

        stdout = b''.join(stdout_parts) if stdout_parts else None
        return Result(stdout, None, exit_code)

    async def _run_ssh(
        self,
        cmd: list[str],
        wd: str | None,
        verbose: bool,
        cmd_input: bytes | None,
        mod_env: dict[str, str] | None,
        interactive: bool,
        log_prefix: str,
    ) -> Result:

        try:
            if interactive:
                if self._has_local_tty():
                    return await self._run_interactive_on_conn(
                        cmd = cmd,
                        wd = wd,
                        cmd_input = cmd_input,
                        mod_env = mod_env,
                    )
                return await self._run_captured_pty_on_conn(
                    cmd = cmd,
                    wd = wd,
                    verbose = verbose,
                    cmd_input = cmd_input,
                    mod_env = mod_env,
                    log_prefix = log_prefix,
                )

            command = self._build_remote_command(cmd, wd)

            stdout_parts: list[bytes] = []
            stderr_parts: list[bytes] = []

            stdout_log_enc = sys.stdout.encoding or 'utf-8'
            stderr_log_enc = sys.stderr.encoding or 'utf-8'

            stdin_mode = asyncssh.PIPE if cmd_input is not None else asyncssh.DEVNULL

            conn = await self._conn

            proc = await conn.create_process(
                command = command,
                env = mod_env,
                stdin = stdin_mode,
                stdout = asyncssh.PIPE,
                stderr = asyncssh.PIPE,
                encoding = None,
                request_pty = False,
            )

            tasks = [
                asyncio.create_task(
                    self._read_stream(
                        proc.stdout,
                        NOTICE,
                        stdout_parts,
                        verbose = verbose,
                        log_prefix = log_prefix,
                        log_enc = stdout_log_enc,
                    )
                ),
                asyncio.create_task(
                    self._read_stream(
                        proc.stderr,
                        ERR,
                        stderr_parts,
                        verbose = verbose,
                        log_prefix = log_prefix,
                        log_enc = stderr_log_enc,
                    )
                ),
            ]

            if cmd_input is not None and proc.stdin is not None:
                proc.stdin.write(cmd_input)
                await proc.stdin.drain()
                proc.stdin.write_eof()

            completed = await proc.wait(check = False)
            await asyncio.gather(*tasks)

            stdout = b''.join(stdout_parts) if stdout_parts else None
            stderr = b''.join(stderr_parts) if stderr_parts else None

            exit_code = completed.exit_status
            if exit_code is None:
                exit_code = (
                    completed.returncode if completed.returncode is not None else -1
                )

            return Result(stdout, stderr, exit_code)

        except Exception as e:
            log(ERR, f'Failed to run command {" ".join(cmd)} ({e})')
            raise