"""SSH command for connecting to running GPU instances.

Provides secure shell access to running tasks for debugging and development.

Examples:
    Connect interactively:
        $ flow ssh task-abc123

    Execute remote command:
        $ flow ssh task-abc123 -c 'nvidia-smi'

    Check GPU utilization:
        $ flow ssh task-abc123 -c 'watch -n1 nvidia-smi'
"""

import sys
from typing import Optional

import click

from flow import Flow
from flow.errors import FlowError

from .base import BaseCommand, console


class SSHCommand(BaseCommand):
    """SSH command implementation.
    
    Handles both interactive sessions and remote command execution.
    Requires task to be in running state with SSH keys configured.
    """

    @property
    def name(self) -> str:
        return "ssh"

    @property
    def help(self) -> str:
        return "SSH to a running task"

    def get_command(self) -> click.Command:
        @click.command(name=self.name, help=self.help)
        @click.argument("task_id")
        @click.option("--command", "-c", help="Command to run on remote host")
        def ssh(task_id: str, command: Optional[str]):
            """SSH to a running task."""
            self._execute(task_id, command)

        return ssh

    def _execute(self, task_id: str, command: Optional[str]) -> None:
        """Execute SSH connection or command.
        
        Args:
            task_id: Task identifier to connect to
            command: Optional command to execute remotely
            
        Raises:
            FlowError: If task not found or SSH unavailable
            SystemExit: On connection failure
        """
        try:
            client = Flow(auto_init=True)

            if command:
                # Run command remotely
                client.ssh(task_id, command=command)
            else:
                # Get task to show SSH info
                task = client.get_task(task_id)

                if not task.ssh_host:
                    console.print(f"[red]✗ Task {task_id} has no SSH access available[/red]")
                    sys.exit(1)

                console.print("\n[bold]Connect with:[/bold]")
                console.print(f"  {task.ssh_command}\n")

                # Actually connect
                client.ssh(task_id)

        except FlowError as e:
            console.print(f"[red]✗ Error:[/red] {e}")
            sys.exit(1)


# Export command instance
command = SSHCommand()
