"""
SSH Connection Management
Handles SSH connection, reconnection, and lifecycle
"""

import paramiko
import socket
import time
import logging
from typing import Optional

logger = logging.getLogger(__name__)


class SSHConnection:
    """Manages SSH connection lifecycle"""

    def __init__(self, host: str = "", user: str = "", password: str = "", port: int = 22,
                 keepalive_interval: int = 30, reconnect_attempts: int = 3,
                 connection_timeout: int = 10):
        """
        Initialize SSH Connection Manager

        Args:
            host: Remote host address (can be empty for multi-server mode)
            user: Username for authentication (can be empty for multi-server mode)
            password: Password for authentication (can be empty for multi-server mode)
            port: SSH port (default 22)
            keepalive_interval: Seconds between keepalive packets
            reconnect_attempts: Number of reconnection attempts
            connection_timeout: Connection timeout in seconds
        """
        self.host = host
        self.user = user
        self.password = password
        self.port = port
        self.keepalive_interval = keepalive_interval
        self.reconnect_attempts = reconnect_attempts
        self.connection_timeout = connection_timeout

        self.client: Optional[paramiko.SSHClient] = None
        self.shell: Optional[paramiko.Channel] = None
        self.connected = False
        self.reconnecting = False
        self._sftp = None

    def reconfigure(self, host: str, user: str, password: str, port: int = 22) -> None:
        """
        Reconfigure connection parameters (disconnect first if connected)

        Args:
            host: New remote host address
            user: New username
            password: New password
            port: New SSH port
        """
        if self.connected:
            self.disconnect()

        self.host = host
        self.user = user
        self.password = password
        self.port = port
        logger.info(f"SSH connection reconfigured for {user}@{host}:{port}")

    def connect(self, host: str = None, user: str = None, password: str = None, port: int = None) -> bool:
        """
        Establish SSH connection

        Args:
            host: Remote host (optional, uses configured if not provided)
            user: Username (optional, uses configured if not provided)
            password: Password (optional, uses configured if not provided)
            port: SSH port (optional, uses configured if not provided)

        Returns:
            True if connection successful, False otherwise
        """
        # Allow override of connection parameters
        if host is not None:
            self.host = host
        if user is not None:
            self.user = user
        if password is not None:
            self.password = password
        if port is not None:
            self.port = port

        try:
            logger.info(f"Connecting to {self.user}@{self.host}:{self.port}")

            ## logger.info(f"DEBUG: password='{self.password}', type={type(self.password)}, len={len(self.password)}")  


            self.client = paramiko.SSHClient()
            self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

            self.client.connect(
                hostname=self.host,
                port=self.port,
                username=self.user,
                password=self.password,
                timeout=self.connection_timeout,
                look_for_keys=False,
                allow_agent=False
            )

            # Set keepalive
            transport = self.client.get_transport()
            if transport:
                transport.set_keepalive(self.keepalive_interval)

            # Get interactive shell
            self.shell = self.client.invoke_shell(
                term='xterm-256color',
                width=120,
                height=40
            )

            self.connected = True
            logger.info("SSH connection established")

            # Wait for initial output (welcome message + prompt)
            time.sleep(0.5)

            return True

        except Exception as e:
            logger.error(f"Connection failed: {e}")
            self.connected = False
            return False

    def disconnect(self) -> None:
        """Close SSH connection"""
        logger.info("Disconnecting SSH")

        # Close SFTP client if exists
        if self._sftp is not None:
            try:
                self._sftp.close()
                logger.info("SFTP client closed")
            except Exception as e:
                logger.warning(f"Error closing SFTP client: {e}")
            finally:
                self._sftp = None

        if self.shell:
            try:
                self.shell.close()
            except:
                pass
            self.shell = None

        if self.client:
            try:
                self.client.close()
            except:
                pass
            self.client = None

        self.connected = False
        logger.info("SSH disconnected")

    def reconnect(self) -> bool:
        """
        Attempt to reconnect to remote machine

        Returns:
            True if reconnection successful
        """
        if self.reconnecting:
            return False

        self.reconnecting = True
        logger.info("Attempting reconnection...")

        self.disconnect()

        for attempt in range(self.reconnect_attempts):
            logger.info(f"Reconnection attempt {attempt + 1}/{self.reconnect_attempts}")
            time.sleep(2 ** attempt)  # Exponential backoff

            if self.connect():
                self.reconnecting = False
                logger.info("Reconnection successful")
                return True

        self.reconnecting = False
        logger.error("Reconnection failed after all attempts")
        return False

    def is_connected(self) -> bool:
        """Check if currently connected"""
        return self.connected and self.client and self.client.get_transport() and \
               self.client.get_transport().is_active()

    def get_connection_info(self) -> dict:
        """Get connection information"""
        return {
            'host': self.host,
            'port': self.port,
            'user': self.user,
            'connected': self.connected,
            'reconnecting': self.reconnecting
        }

    def get_sftp(self) -> paramiko.SFTPClient:
        """
        Get SFTP client from existing SSH connection.
        Creates a new SFTP client if needed or if the existing one is closed.

        Returns:
            paramiko.SFTPClient: Active SFTP client

        Raises:
            RuntimeError: If SSH not connected
        """
        if not self.is_connected():
            raise RuntimeError("SSH not connected. Use connect() first.")

        # Create new SFTP client if needed
        if self._sftp is None or self._sftp.get_channel().closed:
            logger.info("Creating new SFTP client")
            self._sftp = self.client.open_sftp()

        return self._sftp

    def resize_pty(self, cols: int, rows: int) -> bool:
        """
        Resize the pseudo-terminal

        Args:
            cols: Number of columns
            rows: Number of rows

        Returns:
            True if resize successful, False otherwise
        """
        if not self.connected or not self.shell:
            logger.warning("Cannot resize PTY: not connected")
            return False

        try:
            self.shell.resize_pty(width=cols, height=rows)
            logger.debug(f"PTY resized to {cols}x{rows}")
            return True
        except Exception as e:
            logger.error(f"Failed to resize PTY: {e}")
            return False
