import subprocess
import os
import sys

class SSHClient:
    def __init__(self, host, user="root", key_path=None, password=None, timeout=30):
        self.host = host
        self.user = user
        self.key_path = os.path.expanduser(key_path) if key_path else None
        self.password = password
        self.timeout = timeout

    def run(self, cmd, capture=True):
        ssh_cmd = [
            "ssh", 
            "-o", "StrictHostKeyChecking=no", 
            "-o", "UserKnownHostsFile=/dev/null",
            "-o", f"ConnectTimeout={self.timeout}",
            "-o", "BatchMode=yes" if not self.password else "BatchMode=no"
        ]
        
        if self.key_path:
            ssh_cmd += ["-i", self.key_path]
        
        target = f"{self.user}@{self.host}"
        
        if self.password:
            # sshpass is required for password auth
            full_cmd = ["sshpass", "-p", self.password] + ssh_cmd + [target, cmd]
        else:
            full_cmd = ssh_cmd + [target, cmd]

        try:
            result = subprocess.run(
                full_cmd,
                capture_output=capture,
                text=True,
                timeout=self.timeout + 10
            )
            return result
        except subprocess.TimeoutExpired:
            print(f"Error: SSH command timed out after {self.timeout}s on {self.host}", file=sys.stderr)
            # Create a mock result for timeout
            return subprocess.CompletedProcess(full_cmd, 1, "", "Timeout expired")
        except Exception as e:
            print(f"Error: SSH execution failed on {self.host}: {e}", file=sys.stderr)
            return subprocess.CompletedProcess(full_cmd, 1, "", str(e))

    def scp_to(self, local_path, remote_path):
        scp_cmd = ["scp", "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null"]
        if self.key_path:
            scp_cmd += ["-i", self.key_path]
        
        target = f"{self.user}@{self.host}:{remote_path}"
        
        if self.password:
            full_cmd = ["sshpass", "-p", self.password] + scp_cmd + [local_path, target]
        else:
            full_cmd = scp_cmd + [local_path, target]

        return subprocess.run(full_cmd, capture_output=True, timeout=self.timeout + 60)