#!/usr/bin/env python3
"""
ssm -- like SSH, but through AWS Systems Manager.

  ssm myserver                          interactive shell + auto root
  ssm myserver@use1                     shell in us-east-1
  ssm myserver -- uname -a              run a command (SSH-style)
  ssm myserver -c "uname -a"            run a command
  ssm myserver -L 8080                  port forward (same local & remote)
  ssm myserver -L 8080:3000             port forward local 8080 -> remote 3000
  ssm myserver -L 9443:rds.host:443     tunnel via jump EC2
  ssm h1 h2 h3 -c "uptime"              multi-host run
  ssm -f hosts.txt -c "df -h"           multi-host run from file
  ssm myserver --no-sudo                connect without auto root

Accepts Name, FQDN, private/public IP, or instance-id as TARGET.
Region: --region/-R (use1|euw1|aps1 or full), $AWS_REGION, or us-east-1.
Mode (readonly/devops/admin) auto-detected from your Duo/SSO role.
Auto Duo SSO re-login when token is expired (override: --duo-cmd / $SSM_DUO_CMD).
Auto sudo su - on admin/devops shells (disable: --no-sudo).
"""
import argparse
import base64
import fcntl
import json
import os
import pty
import re
import select
import shlex
import signal
import subprocess
import sys
import termios
import threading
import time
import tty
from typing import List, Optional, Tuple

import boto3
import botocore.session
from botocore import credentials
from botocore.exceptions import ClientError, ProfileNotFound

# ── Constants ─────────────────────────────────────────────────────────────────

REGIONS = {
    "use1": "us-east-1",
    "euw1": "eu-west-1",
    "aps1": "ap-southeast-1",
}

EXPIRED_TOKEN_PATTERNS = re.compile(
    r"ExpiredToken|expired token|security token.*expired|Token has expired|"
    r"SSO.*expired|sso.*expired|SSO session.*expired|InvalidClientTokenId|"
    r"UnrecognizedClientException",
    re.IGNORECASE,
)
INSTANCE_ID_PATTERN = re.compile(r"^i-[0-9a-fA-F]{8,17}$")
IPV4_PATTERN = re.compile(r"^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$")
READONLY_ROLE_PATTERN = re.compile(
    r"(^|[^a-z0-9])(ssm[-_]?readonly|read[-_]?only)([^a-z0-9]|$)", re.IGNORECASE
)
DEVOPS_ROLE_PATTERN = re.compile(r"(^|[^a-z0-9])devops([^a-z0-9]|$)", re.IGNORECASE)
ASSUMED_ROLE_PATTERN = re.compile(r":assumed-role/([^/]+)/")

DEFAULT_DUO_COMMAND = "duo-sso --session-duration 3600"
ANSI_ESCAPE_RE = re.compile(rb"\x1b\[[^a-zA-Z]*[a-zA-Z]")


# ── UI helpers ────────────────────────────────────────────────────────────────

_IS_TTY = hasattr(sys.stderr, "isatty") and sys.stderr.isatty()


def _ansi(code: str, text: str) -> str:
    """Wrap text in ANSI escape code when stderr is a TTY."""
    return f"\033[{code}m{text}\033[0m" if _IS_TTY else str(text)


def _dim(text: str) -> str:
    return _ansi("2", text)


def _bold(text: str) -> str:
    return _ansi("1", text)


def _green(text: str) -> str:
    return _ansi("32", text)


def _red(text: str) -> str:
    return _ansi("31", text)


def _cyan(text: str) -> str:
    return _ansi("36", text)


class Spinner:
    """Braille spinner on TTY; falls back to a static message otherwise."""

    _FRAMES = "⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏"

    def __init__(self, message: str):
        self.message = message
        self._stop_event = threading.Event()
        self._thread: Optional[threading.Thread] = None

    def __enter__(self):
        if _IS_TTY:
            self._thread = threading.Thread(target=self._spin, daemon=True)
            self._thread.start()
        return self

    def __exit__(self, *_):
        self._stop_event.set()
        if self._thread:
            self._thread.join()
        if _IS_TTY:
            sys.stderr.write("\r\033[K")
            sys.stderr.flush()

    def _spin(self):
        idx = 0
        while not self._stop_event.wait(0.08):
            frame = self._FRAMES[idx % len(self._FRAMES)]
            sys.stderr.write(f"\r  {frame} {self.message}")
            sys.stderr.flush()
            idx += 1


def _info(message: str) -> None:
    sys.stderr.write(f"  {message}\n")
    sys.stderr.flush()


def _ok(message: str) -> None:
    sys.stderr.write(f"  {_green('✓')} {message}\n")
    sys.stderr.flush()


def _fail(message: str) -> None:
    sys.stderr.write(f"  {_red('✗')} {message}\n")
    sys.stderr.flush()


# ── Core logic ────────────────────────────────────────────────────────────────


def _make_session(profile: str) -> boto3.Session:
    """Create boto3 Session with assume-role credential cache."""
    cli_cache = os.path.join(os.path.expanduser("~"), ".aws/cli/cache")
    botocore_session = botocore.session.Session(profile=profile)
    botocore_session.get_component("credential_provider").get_provider(
        "assume-role"
    ).cache = credentials.JSONFileCache(cli_cache)
    return boto3.Session(botocore_session=botocore_session, profile_name=profile)


def _check_auth_status(session: boto3.Session, region: str) -> Tuple[bool, bool, str]:
    """Check AWS credentials.

    Returns (is_ok, is_token_expired, error_message).
    """
    sts = session.client("sts", region_name=region)
    try:
        sts.get_caller_identity()
        return True, False, ""
    except ClientError as exc:
        msg = (exc.response.get("Error", {}).get("Message") or str(exc)) or ""
        is_expired = bool(EXPIRED_TOKEN_PATTERNS.search(msg))
        return False, is_expired, msg


def check_auth(session: boto3.Session, region: str) -> bool:
    """Verify AWS identity; print clear message if token expired."""
    is_ok, is_expired, msg = _check_auth_status(session, region)
    if is_ok:
        return True
    if is_expired:
        _fail("AWS/Duo token expired. Please re-authenticate and retry.")
        _info("Tip: run your Duo/AWS SSO login (e.g. 'duo-sso'), then re-run.")
    else:
        _fail(f"AWS authentication failed: {msg}")
    return False


def ensure_auth(
    session: boto3.Session,
    profile: str,
    region: str,
    duo_command: Optional[str] = None,
) -> boto3.Session:
    """Check auth; if the token is expired, run Duo SSO automatically and retry.

    Returns a valid boto3 Session or exits with code 1.
    """
    duo_cmd = duo_command or os.environ.get("SSM_DUO_CMD") or DEFAULT_DUO_COMMAND

    is_ok, is_expired, error_msg = _check_auth_status(session, region)
    if is_ok:
        return session

    if not is_expired:
        _fail(f"AWS authentication failed: {error_msg}")
        sys.exit(1)

    # Token expired — auto re-login via Duo SSO
    _info(f"Token expired — running Duo SSO re-login ...")
    _info(f"  $ {_dim(duo_cmd)}")
    cmd_parts = shlex.split(duo_cmd) if isinstance(duo_cmd, str) else list(duo_cmd)
    returncode = subprocess.call(cmd_parts)
    if returncode != 0:
        _fail("Duo SSO re-login failed.")
        sys.exit(1)

    # Recreate session with fresh credentials
    session = _make_session(profile)
    is_ok, _, error_msg = _check_auth_status(session, region)
    if is_ok:
        _ok("Re-authenticated successfully")
        return session

    _fail(f"Still not authenticated after Duo SSO: {error_msg}")
    sys.exit(1)


def resolve_instance_id(
    session: boto3.Session, identifier: str, region: str
) -> Optional[str]:
    """Resolve EC2 identifier (instance-id / IP / Name / FQDN) to instance-id."""
    if INSTANCE_ID_PATTERN.match(identifier):
        return identifier

    ec2 = session.client("ec2", region_name=region)

    if IPV4_PATTERN.match(identifier):
        for filter_name in ("private-ip-address", "ip-address"):
            try:
                resp = ec2.describe_instances(
                    Filters=[
                        {"Name": filter_name, "Values": [identifier]},
                        {"Name": "instance-state-name", "Values": ["running"]},
                    ],
                )
                instances = [
                    i["InstanceId"]
                    for r in resp.get("Reservations", [])
                    for i in r.get("Instances", [])
                ]
                if instances:
                    return instances[0]
            except ClientError:
                continue
        return None

    short_name = identifier.split(".")[0]
    try:
        resp = ec2.describe_instances(
            Filters=[
                {"Name": "tag:Name", "Values": [short_name]},
                {"Name": "instance-state-name", "Values": ["running"]},
            ],
        )
        instances = [
            i["InstanceId"]
            for r in resp.get("Reservations", [])
            for i in r.get("Instances", [])
        ]
        return instances[0] if instances else None
    except ClientError:
        return None


def detect_mode(session: boto3.Session, region: str) -> str:
    """Detect readonly / devops / admin from SSM_MODE env or caller ARN."""
    mode = os.environ.get("SSM_MODE", "").strip().lower()
    if mode in ("readonly", "devops", "admin"):
        return mode
    try:
        sts = session.client("sts", region_name=region)
        identity = sts.get_caller_identity()
        arn = identity.get("Arn") or ""
    except ClientError:
        return "admin"
    if not arn or arn == "None":
        return "admin"
    match = ASSUMED_ROLE_PATTERN.search(arn)
    role = match.group(1) if match else ""
    check = (role or arn).lower()
    if READONLY_ROLE_PATTERN.search(check):
        return "readonly"
    if DEVOPS_ROLE_PATTERN.search(check):
        return "devops"
    return "admin"


# ── Multi-host helpers ────────────────────────────────────────────────────────


def collect_targets_from_file(path: str) -> List[str]:
    """Read host identifiers from file; strip comments and blank lines."""
    with open(path) as fh:
        return [
            line.strip()
            for line in fh
            if line.strip() and not line.strip().startswith("#")
        ]


def _chunks(items: List[str], size: int):
    """Yield successive chunks of *items* with at most *size* elements."""
    for i in range(0, len(items), size):
        yield items[i : i + size]


def build_ident_to_instance_id_map(
    session: boto3.Session, region: str, identifiers: List[str]
) -> List[Tuple[str, str]]:
    """Batch-resolve identifiers to instance-ids. Returns [(ident, iid), ...]."""
    ec2 = session.client("ec2", region_name=region)
    result: List[Tuple[str, str]] = []

    direct_ids = [x for x in identifiers if INSTANCE_ID_PATTERN.match(x)]
    for ident in direct_ids:
        result.append((ident, ident))

    ip_idents = [
        x for x in identifiers if IPV4_PATTERN.match(x) and x not in direct_ids
    ]
    name_idents = [
        x for x in identifiers if x not in direct_ids and not IPV4_PATTERN.match(x)
    ]

    ip_to_iid: dict = {}
    for filter_name, ip_key in [
        ("private-ip-address", "PrivateIpAddress"),
        ("ip-address", "PublicIpAddress"),
    ]:
        for batch in _chunks(ip_idents, 50):
            try:
                resp = ec2.describe_instances(
                    Filters=[
                        {"Name": filter_name, "Values": batch},
                        {"Name": "instance-state-name", "Values": ["running"]},
                    ],
                )
                for r in resp.get("Reservations", []):
                    for i in r.get("Instances", []):
                        ip_val = i.get(ip_key)
                        if ip_val and ip_val in ip_idents and ip_val not in ip_to_iid:
                            ip_to_iid[ip_val] = i["InstanceId"]
            except ClientError:
                pass
    for ip_val in ip_idents:
        if ip_val in ip_to_iid:
            result.append((ip_val, ip_to_iid[ip_val]))

    for ident in name_idents:
        iid = resolve_instance_id(session, ident, region)
        if iid:
            result.append((ident, iid))

    return result


def get_online_instance_ids(
    ssm_client, instance_ids: List[str], chunk_size: int = 50
) -> List[str]:
    """Return instance IDs that are SSM Online (batched)."""
    online: List[str] = []
    for batch in _chunks(instance_ids, chunk_size):
        try:
            resp = ssm_client.describe_instance_information(
                Filters=[{"Key": "InstanceIds", "Values": batch}],
            )
            for info in resp.get("InstanceInformationList", []):
                if info.get("PingStatus") == "Online":
                    online.append(info["InstanceId"])
        except ClientError:
            pass
    return online


def run_command(
    session: boto3.Session,
    region: str,
    profile: str,
    identifiers: List[str],
    command: Optional[str] = None,
    script_path: Optional[str] = None,
    timeout_seconds: Optional[int] = None,
) -> bool:
    """Run command / script on resolved, SSM-online instances.  Returns True on full success."""
    timeout_seconds = timeout_seconds or int(
        os.environ.get("SSM_TIMEOUT_SECONDS", "3600")
    )

    with Spinner(f"Resolving {len(identifiers)} target(s) in {region}"):
        ident_to_iid = build_ident_to_instance_id_map(session, region, identifiers)
    all_iids = list({iid for _, iid in ident_to_iid})
    if not all_iids:
        _fail("No instance-ids resolved. Aborting.")
        return False

    ssm_client = session.client("ssm", region_name=region)
    with Spinner("Checking SSM agent status"):
        online_ids = get_online_instance_ids(ssm_client, all_iids)

    ident_by_iid = {iid: ident for ident, iid in ident_to_iid}
    skipped = [iid for iid in all_iids if iid not in online_ids]
    if skipped:
        for iid in skipped:
            _info(
                f"  ⚠  {ident_by_iid.get(iid, iid)} → {iid}  (not SSM Online, skipped)"
            )
    if not online_ids:
        _fail("None of the targets are SSM Online.")
        return False

    _ok(f"{len(online_ids)} target(s) online")

    # Build the command payload
    if script_path:
        with open(script_path) as fh:
            script_b64 = base64.b64encode(fh.read().encode()).decode()
        remote_cmd = (
            f"TMPFILE=/tmp/ssm-run-{int(time.time())}-$$.sh; "
            f"printf '%s' '{script_b64}' | base64 -d 2>/dev/null || base64 --decode > \"$TMPFILE\"; "
            f'chmod +x "$TMPFILE"; bash "$TMPFILE"; rc=$?; rm -f "$TMPFILE"; exit $rc'
        )
        commands = [f'bash -lc "{remote_cmd}"']
    else:
        if not command:
            _fail('Missing command. Use -c "command" or -- command.')
            return False
        commands = [f'bash -lc "{command}"']

    try:
        send_resp = ssm_client.send_command(
            InstanceIds=online_ids,
            DocumentName="AWS-RunShellScript",
            Comment=f"ssm-connect run (profile={profile})",
            TimeoutSeconds=timeout_seconds,
            Parameters={"commands": commands},
        )
    except ClientError as exc:
        _fail(f"Failed to send SSM command: {exc}")
        return False

    cmd_id = send_resp["Command"]["CommandId"]
    _info(f"CommandId: {_dim(cmd_id)}")

    # Wait for results and display per-host output
    ok_count, fail_count, other_count = 0, 0, 0
    for iid in online_ids:
        ident = ident_by_iid.get(iid, iid)
        with Spinner(f"Waiting for {ident}"):
            for _ in range(120):
                try:
                    inv = ssm_client.get_command_invocation(
                        CommandId=cmd_id, InstanceId=iid
                    )
                    if inv.get("Status") in (
                        "Success",
                        "Failed",
                        "Cancelled",
                        "TimedOut",
                    ):
                        break
                except ClientError:
                    pass
                time.sleep(1)

        inv = ssm_client.get_command_invocation(CommandId=cmd_id, InstanceId=iid)
        status = inv.get("Status", "UNKNOWN")
        stdout = inv.get("StandardOutputContent") or ""
        stderr = inv.get("StandardErrorContent") or ""

        if status == "Success":
            ok_count += 1
            status_icon = _green("✓")
        elif status == "Failed":
            fail_count += 1
            status_icon = _red("✗")
        else:
            other_count += 1
            status_icon = _dim("?")

        print(f"  {status_icon} {_bold(ident)} {_dim(f'({iid})')}")
        if stdout.strip():
            for line in stdout.strip().splitlines():
                print(f"    {line}")
        if stderr.strip():
            for line in stderr.strip().splitlines():
                print(f"    {_red(line)}")

    print()
    _info(
        f"Done: {_green(str(ok_count))} ok"
        f", {_red(str(fail_count))} failed"
        f", {_dim(str(other_count))} other"
    )
    return fail_count == 0 and other_count == 0


# ── CLI helpers ───────────────────────────────────────────────────────────────


def _normalize_region(region: str) -> str:
    """Resolve shortcut (use1, euw1, aps1) to full AWS region name."""
    return REGIONS.get(region.lower(), region)


def _parse_target_at_region(raw: str) -> Tuple[str, Optional[str]]:
    """Parse 'identifier@region' -> (identifier, region_or_None)."""
    if "@" in raw and not INSTANCE_ID_PATTERN.match(raw):
        identifier, region_hint = raw.rsplit("@", 1)
        return identifier, region_hint
    return raw, None


def _parse_forward_spec(spec: str) -> Tuple[Optional[str], str, str]:
    """Parse SSH-style -L spec.

    Returns (remote_host_or_None, remote_port, local_port).

      -L 8080           -> (None,   "8080", "8080")   direct PF same port
      -L 8080:3000      -> (None,   "3000", "8080")   direct PF local:remote
      -L 9443:host:443  -> ("host", "443",  "9443")   tunnel local:host:remote
    """
    parts = spec.split(":")
    if len(parts) == 1:
        return None, parts[0], parts[0]
    if len(parts) == 2:
        return None, parts[1], parts[0]
    # 3+ parts: local:host:remote  (host may contain dots but not colons in practice)
    return parts[1], parts[2], parts[0]


# ── Auto-root (pty-based) ────────────────────────────────────────────────────


def _copy_terminal_size(source_fd: int, target_fd: int) -> None:
    """Copy terminal window dimensions from *source_fd* to *target_fd*."""
    try:
        winsize = fcntl.ioctl(source_fd, termios.TIOCGWINSZ, b"\x00" * 8)
        fcntl.ioctl(target_fd, termios.TIOCSWINSZ, winsize)
    except OSError:
        pass


def _spawn_with_auto_root(aws_cmd: list) -> int:
    """Spawn *aws_cmd* in a pty, auto-send ``sudo su -`` on first shell prompt.

    * Detects ``$ `` / ``% `` prompt → sends ``sudo su -``
    * Detects ``# `` prompt → already root, skips escalation
    * 30-second timeout → gives up silently if no prompt detected
    * Propagates SIGWINCH so terminal resize works mid-session
    * Returns the child process exit code.
    """
    pid, master_fd = pty.fork()
    if pid == 0:
        # Child — becomes the SSM session
        os.execvp(aws_cmd[0], aws_cmd)
        os._exit(1)  # unreachable unless execvp fails

    # Parent — relay I/O between user's terminal and the pty
    stdin_fd = sys.stdin.fileno()
    stdout_fd = sys.stdout.fileno()
    _copy_terminal_size(stdin_fd, master_fd)

    # Propagate terminal resize
    prev_winch = signal.getsignal(signal.SIGWINCH)

    def _on_winch(_sig, _frame):
        _copy_terminal_size(stdin_fd, master_fd)

    signal.signal(signal.SIGWINCH, _on_winch)

    # Switch to raw mode so keystrokes pass through immediately
    try:
        old_tty_attrs = termios.tcgetattr(stdin_fd)
        tty.setraw(stdin_fd)
        restore_tty = True
    except termios.error:
        restore_tty = False

    escalated = False
    output_buffer = b""
    prompt_re = re.compile(rb"[\$%]\s*$")
    root_prompt_re = re.compile(rb"#\s*$")
    escalation_deadline = time.monotonic() + 30

    try:
        while True:
            try:
                readable, _, _ = select.select([master_fd, stdin_fd], [], [], 0.25)
            except InterruptedError:
                continue  # SIGWINCH can interrupt select

            if master_fd in readable:
                try:
                    data = os.read(master_fd, 4096)
                except OSError:
                    break
                if not data:
                    break
                os.write(stdout_fd, data)

                # Prompt detection for auto-root
                if not escalated:
                    if time.monotonic() > escalation_deadline:
                        escalated = True  # timeout — give up silently
                    else:
                        output_buffer += data
                        if len(output_buffer) > 8192:
                            output_buffer = output_buffer[-4096:]
                        clean = ANSI_ESCAPE_RE.sub(b"", output_buffer)
                        if root_prompt_re.search(clean):
                            escalated = True  # already root
                        elif prompt_re.search(clean):
                            time.sleep(0.15)
                            os.write(master_fd, b"sudo su -\n")
                            escalated = True

            if stdin_fd in readable:
                try:
                    data = os.read(stdin_fd, 4096)
                except OSError:
                    break
                if not data:
                    break
                os.write(master_fd, data)
    except OSError:
        pass
    finally:
        if restore_tty:
            termios.tcsetattr(stdin_fd, termios.TCSAFLUSH, old_tty_attrs)
        signal.signal(signal.SIGWINCH, prev_winch)

    os.close(master_fd)
    _, status = os.waitpid(pid, 0)
    return os.WEXITSTATUS(status) if os.WIFEXITED(status) else 1


# ── Action handlers ──────────────────────────────────────────────────────────


def _handle_shell(
    session: boto3.Session,
    profile: str,
    region: str,
    identifier: str,
    auto_root: bool = True,
) -> None:
    """Open an interactive SSM shell — like 'ssh hostname'.

    When *auto_root* is True, ``sudo su -`` is sent automatically after
    the first shell prompt appears (works in all modes).
    """
    start = time.monotonic()
    with Spinner(f"Resolving {_bold(identifier)} in {region}"):
        instance_id = resolve_instance_id(session, identifier, region)
    elapsed = time.monotonic() - start

    if not instance_id:
        _fail(f"No running instance for '{identifier}' in {region}")
        sys.exit(1)

    display = (
        identifier if instance_id == identifier else f"{identifier} → {instance_id}"
    )
    _ok(f"{display}  {_dim(f'{elapsed:.1f}s')}")

    mode = detect_mode(session, region)
    mode_labels = {
        "readonly": "read-only 🔒",
        "devops": "devops 🔧",
        "admin": "admin 🔑",
    }
    _ok(f"Mode: {mode_labels.get(mode, mode)}")

    document = "SSM-SessionManagerReadOnlyShell" if mode == "readonly" else None
    aws_cmd = [
        "aws",
        "ssm",
        "start-session",
        "--target",
        instance_id,
        "--profile",
        profile,
        "--region",
        region,
    ]
    if document:
        aws_cmd.extend(["--document-name", document])

    will_auto_root = auto_root and sys.stdin.isatty()
    if will_auto_root:
        _ok("Auto-root enabled (sudo su -)")

    _info(f"Connecting to {_bold(identifier)} ...")
    print()

    if will_auto_root:
        exit_code = _spawn_with_auto_root(aws_cmd)
        sys.exit(exit_code)
    else:
        sys.exit(subprocess.call(aws_cmd))


def _handle_forward(
    session: boto3.Session,
    profile: str,
    region: str,
    identifier: str,
    forward_spec: str,
) -> None:
    """Direct port forward or tunnel via jump EC2 (decided by -L spec)."""
    remote_host, remote_port, local_port = _parse_forward_spec(forward_spec)

    start = time.monotonic()
    with Spinner(f"Resolving {_bold(identifier)} in {region}"):
        instance_id = resolve_instance_id(session, identifier, region)
    elapsed = time.monotonic() - start

    if not instance_id:
        _fail(f"No running instance for '{identifier}' in {region}")
        sys.exit(1)

    _ok(f"{identifier} → {instance_id}  {_dim(f'{elapsed:.1f}s')}")

    if remote_host:
        # Tunnel mode: local -> jump EC2 -> remote_host:remote_port
        _info(
            f"Tunnel: localhost:{_bold(local_port)}"
            f" → {_bold(remote_host)}:{_bold(remote_port)}"
            f"  via {_dim(instance_id)}"
        )
        params = json.dumps(
            {
                "host": [remote_host],
                "portNumber": [remote_port],
                "localPortNumber": [local_port],
            }
        )
        document = "AWS-StartPortForwardingSessionToRemoteHost"
    else:
        # Direct port forward to the EC2 instance
        _info(
            f"Port forward: localhost:{_bold(local_port)}"
            f" → {_bold(identifier)}:{_bold(remote_port)}"
        )
        params = json.dumps(
            {
                "portNumber": [remote_port],
                "localPortNumber": [local_port],
            }
        )
        document = "AWS-StartPortForwardingSession"

    print()
    sys.exit(
        subprocess.call(
            [
                "aws",
                "ssm",
                "start-session",
                "--target",
                instance_id,
                "--document-name",
                document,
                "--parameters",
                params,
                "--profile",
                profile,
                "--region",
                region,
            ]
        )
    )


def _handle_run(
    session: boto3.Session,
    profile: str,
    region: str,
    targets: List[str],
    command: Optional[str],
    script_file: Optional[str],
) -> None:
    """Run command or script on one or many targets."""
    if script_file and not os.path.isfile(script_file):
        _fail(f"Script file not found: {script_file}")
        sys.exit(1)
    success = run_command(
        session,
        region,
        profile,
        targets,
        command=command,
        script_path=script_file,
    )
    sys.exit(0 if success else 1)


# ── Argument parser ──────────────────────────────────────────────────────────

_EXAMPLES = """\
examples:
  %(prog)s myserver                          interactive shell
  %(prog)s myserver@use1                     shell in us-east-1
  %(prog)s myserver -- uname -a              run a command (SSH-style)
  %(prog)s myserver -c "uname -a"            run a command
  %(prog)s myserver -L 8080                  port forward (same local & remote)
  %(prog)s myserver -L 8080:3000             port forward local 8080 -> remote 3000
  %(prog)s myserver -L 9443:rds.host:443     tunnel via jump EC2
  %(prog)s h1 h2 h3 -c "uptime"             multi-host run
  %(prog)s -f hosts.txt -c "df -h"           multi-host run from file"""


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="ssm",
        description="ssm — like SSH, but through AWS Systems Manager.",
        epilog=_EXAMPLES,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "targets",
        nargs="*",
        metavar="TARGET",
        help="identifier[@region]  (Name, FQDN, IP, or instance-id)",
    )
    parser.add_argument(
        "-p",
        "--profile",
        default="default",
        metavar="PROFILE",
        help="AWS profile (default: default)",
    )
    parser.add_argument(
        "-R",
        "--region",
        metavar="REGION",
        help="AWS region or shortcut (use1 | euw1 | aps1)",
    )
    parser.add_argument(
        "-L",
        dest="forward",
        metavar="[L:]R | L:HOST:R",
        help="port forward / tunnel  (like ssh -L)",
    )
    parser.add_argument(
        "-c",
        "--command",
        metavar="CMD",
        help="run CMD on target(s) instead of opening a shell",
    )
    parser.add_argument(
        "-f",
        "--file",
        dest="hosts_file",
        metavar="FILE",
        help="read targets from FILE (one per line)",
    )
    parser.add_argument(
        "--script",
        dest="script_file",
        metavar="SCRIPT",
        help="upload and execute SCRIPT on target(s)",
    )
    parser.add_argument(
        "--no-sudo",
        dest="no_sudo",
        action="store_true",
        help="disable automatic root escalation (sudo su -)",
    )
    parser.add_argument(
        "--duo-cmd",
        dest="duo_cmd",
        metavar="CMD",
        help=f"Duo SSO re-login command (default: {DEFAULT_DUO_COMMAND})",
    )
    return parser


# ── Entry point ──────────────────────────────────────────────────────────────


def main() -> None:
    # ── Split at '--' for SSH-style command passing: ssm host -- cmd args ──
    argv = sys.argv[1:]
    trailing_command: Optional[str] = None
    if "--" in argv:
        sep = argv.index("--")
        trailing_command = " ".join(argv[sep + 1 :])
        argv = argv[:sep]

    parser = _build_parser()
    args = parser.parse_args(argv)

    # '--' command takes precedence over -c
    if trailing_command is not None:
        args.command = trailing_command

    # ── Collect targets ───────────────────────────────────────────────────
    targets: List[str] = list(args.targets or [])
    if args.hosts_file:
        try:
            targets.extend(collect_targets_from_file(args.hosts_file))
        except FileNotFoundError:
            _fail(f"Hosts file not found: {args.hosts_file}")
            sys.exit(1)

    if not targets:
        parser.print_help()
        sys.exit(0)

    # ── Parse target@region ───────────────────────────────────────────────
    explicit_region = args.region
    parsed_targets: List[str] = []
    for raw in targets:
        ident, region_hint = _parse_target_at_region(raw)
        parsed_targets.append(ident)
        if region_hint and not explicit_region:
            explicit_region = region_hint

    region = _normalize_region(
        explicit_region
        or os.environ.get("AWS_REGION")
        or os.environ.get("AWS_DEFAULT_REGION")
        or "us-east-1"
    )
    profile = args.profile

    # ── Authenticate (auto Duo SSO re-login if expired) ───────────────────
    try:
        session = _make_session(profile)
    except ProfileNotFound:
        _fail(f"AWS profile '{profile}' not found. Check ~/.aws/config.")
        sys.exit(1)

    session = ensure_auth(
        session,
        profile,
        region,
        duo_command=getattr(args, "duo_cmd", None),
    )

    # ── Dispatch ──────────────────────────────────────────────────────────
    auto_root = not getattr(args, "no_sudo", False)

    if args.forward:
        if not parsed_targets:
            _fail("Port forwarding requires a TARGET.")
            sys.exit(1)
        _handle_forward(session, profile, region, parsed_targets[0], args.forward)

    elif args.command or args.script_file:
        if not parsed_targets:
            _fail("Running a command requires at least one TARGET.")
            sys.exit(1)
        _handle_run(
            session,
            profile,
            region,
            parsed_targets,
            args.command,
            args.script_file,
        )

    elif len(parsed_targets) == 1:
        _handle_shell(session, profile, region, parsed_targets[0], auto_root=auto_root)

    elif len(parsed_targets) > 1:
        _fail("Multiple targets require -c CMD or --script FILE.")
        sys.exit(1)

    else:
        parser.print_help()


def _sigint_handler(_signum, _frame):
    """Suppress traceback on Ctrl+C; let the SSM subprocess handle it."""
    pass


def entry_point() -> None:
    """Console script entry point (called by ``ssm`` command after pip install)."""
    signal.signal(signal.SIGINT, _sigint_handler)
    main()


if __name__ == "__main__":
    entry_point()
