from typing import IO

from .Credentials import CredentialCollection
from ..Logging.Logger import log, Severity
import paramiko

class SSHClient:
    def __init__(self, credentials: CredentialCollection, host: str, port: int=22):
        self.credential_collection = credentials
        self.client = None
        self.host = host
        self.port = port

    def connect(self, authorize: bool = True):
        """
        Connect to the SSH server using the provided credentials.
        Falls back to username/password authentication if private key authentication fails.

        Args:
            authorize: Whether to authorize the local machine to the host if necessary.

        Raises:
            paramiko.AuthenticationException: If authentication fails with both private key and password.
            paramiko.SSHException: If there is an error establishing the SSH connection.
        """
        self.client = paramiko.SSHClient()
        self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())  # Add new fingerprints to known_hosts

        # Try to use keyfile first (preferred)
        keyfile = self.credential_collection.get_private_keyfile()
        username = self.credential_collection.username
        password = self.credential_collection.password
        if keyfile:
            private_key = paramiko.RSAKey.from_private_key_file(keyfile)
            try:
                self.client.connect(self.host, port=self.port, username=username, pkey=private_key)
            except paramiko.AuthenticationException:
                # Fallback to password if private key authentication fails
                self.client.connect(self.host, port=self.port, username=username, password=password)

                if authorize:
                    self.add_public_key_to_authorized_keys()
        else:
            # Directly use password if no keyfile is available
            self.client.connect(self.host, port=self.port, username=username, password=password)

    def execute_command(self, command: str, non_blocking=False) -> tuple[str, str, int]:
        """
        Execute a command on the remote server and return the output and result code.

        Args:
        command: The command to execute on the remote server.

        Returns:
            tuple[str, str, int]: A tuple containing:
                - stdout: The standard output of the command as a string.
                - stderr: The standard error of the command as a string.
                - result_code: The exit status code of the command.

        Raises:
            RuntimeError: If the SSH client is not connected
        """
        if not self.client:
            raise ConnectionError("SSH client is not connected. Call connect() first.")

        stdin, stdout, stderr = self.client.exec_command(command)
        if non_blocking:
            return '', '', 0
        output = stdout.read().decode()
        error = stderr.read().decode()
        result_code = stdout.channel.recv_exit_status()

        return output, error, result_code

    def execute_command_live(self, command: str) ->  tuple[IO, IO, IO]:
        if not self.client:
            raise ConnectionError("SSH client is not connected. Call connect() first.")

        stdin, stdout, stderr = self.client.exec_command(command)

        return stdin, stdout, stderr

    def close(self):
        """Close the SSH connection."""
        if self.client:
            self.client.close()
            self.client = None

    def read_file(self, remote_path: str) -> bytes:
        """
        Read a file from the remote SSH server and return its contents as bytes.

        Args:
            remote_path: Path to the file on the remote server

        Returns:
            bytes: Contents of the file

        Raises:
            ConnectionError: If the SSH client is not connected
            RuntimeError: If the file cannot be read
        """
        if not self.client:
            raise ConnectionError("SSH client is not connected. Call connect() first.")

        try:
            # Use SFTP to transfer the file
            sftp = self.client.open_sftp()
            with sftp.open(remote_path, 'rb') as remote_file:
                return remote_file.read()
        except IOError as e:
            raise RuntimeError(f"Failed to read file {remote_path}: {str(e)}")
        finally:
            if 'sftp' in locals():
                # noinspection PyUnboundLocalVariable
                sftp.close()

    def write_file(self, remote_path: str, content: bytes) -> None:
        """
        Write bytes content to a file on the remote SSH server.

        Args:
            remote_path: Path to the file on the remote server
            content: Bytes content to write to the file

        Raises:
            ConnectionError: If the SSH client is not connected
            RuntimeError: If the file cannot be written
        """
        if not self.client:
            raise ConnectionError("SSH client is not connected. Call connect() first.")

        try:
            # Use SFTP to transfer the file
            sftp = self.client.open_sftp()
            with sftp.file(remote_path, 'wb') as remote_file:
                remote_file.write(content)
        except IOError as e:
            raise RuntimeError(f"Failed to write file {remote_path}: {str(e)}")
        finally:
            if 'sftp' in locals():
                # noinspection PyUnboundLocalVariable
                sftp.close()

    def file_exists(self, remote_path: str) -> bool:
        """
        Check if a file exists at the given path on the remote SSH server.

        Args:
            remote_path: Path to the file on the remote server

        Returns:
            bool: True if the file exists, False otherwise

        Raises:
            ConnectionError: If the SSH client is not connected
        """
        if not self.client:
            raise ConnectionError("SSH client is not connected. Call connect() first.")

        try:
            sftp = self.client.open_sftp()
            try:
                sftp.stat(remote_path)
                return True
            except IOError:
                return False
        except IOError as e:
            raise RuntimeError(f"Failed to check file existence: {str(e)}")
        finally:
            if 'sftp' in locals():
                sftp.close()

    def add_public_key_to_authorized_keys(self) -> None:
        """
        Add a public key to the remote machine's authorized_keys file using password authentication.

        Raises:
            ConnectionError: If the SSH client is not connected
            RuntimeError: If the file cannot be written
        """

        stdout, stderr, retval = self.execute_command('sh -c "cd ~ && pwd"')
        if retval != 0:
            raise RuntimeError(f"Failed to get home directory: {stderr}")
        user_homedir = stdout.strip()
        log(f"Adding public key to authorized keys for {self.host} to {user_homedir}", Severity.TRACE)

        # Read the public key from the file specified in the credential collection l
        with open(self.credential_collection.get_public_keyfile(), "r") as key_file:
            public_key = key_file.read()

        self.execute_command(f"mkdir {user_homedir}/.ssh")
        self.execute_command(f"chmod 700 {user_homedir}/.ssh")
        if not self.client:
            raise ConnectionError("SSH client is not connected. Call connect() first.")

        try:
            # Use SFTP to transfer the file
            sftp = self.client.open_sftp()
            with sftp.file(f"{user_homedir}/.ssh/authorized_keys", 'a') as remote_file:
                remote_file.write(public_key.encode() + b'\n')
        except IOError as e:
            raise RuntimeError(f"Failed to write to authorized_keys file: {str(e)}")
        finally:
            if 'sftp' in locals():
                # noinspection PyUnboundLocalVariable
                sftp.close()

        self.execute_command(f"chmod 600 {user_homedir}/.ssh/authorized_keys")
