from __future__ import annotations

import argparse
import getpass
import json
import os
import shlex
import shutil
import subprocess
import sys
import uuid
from pathlib import Path
from typing import Any, Optional, Sequence

from ._cli_config import (
    DEFAULT_BILLING_URL,
    DEFAULT_DOCS_URL,
    ConfigError,
    get_active_profile,
    redact_secret,
    resolve_runtime_config,
    update_profile_settings,
)
from .exceptions import AuthenticationError, InstaVMError
from .sandbox_client import InstaVM


def _print_json(payload: Any) -> None:
    sys.stdout.write(json.dumps(payload) + "\n")


def _print_lines(*lines: str) -> None:
    for line in lines:
        sys.stdout.write(line + "\n")


def _require_api_key(args: argparse.Namespace) -> tuple[InstaVM, Any]:
    settings = resolve_runtime_config(
        api_key=getattr(args, "api_key", None),
        base_url=getattr(args, "base_url", None),
        ssh_host=getattr(args, "ssh_host", None),
    )
    if not settings.api_key:
        raise AuthenticationError(
            "No InstaVM API key configured. Run `instavm auth set-key` or export INSTAVM_API_KEY."
        )
    return InstaVM(
        api_key=settings.api_key,
        base_url=settings.base_url,
        auto_start_session=False,
    ), settings


def _resolve_settings(args: argparse.Namespace):
    return resolve_runtime_config(
        api_key=getattr(args, "api_key", None),
        base_url=getattr(args, "base_url", None),
        ssh_host=getattr(args, "ssh_host", None),
    )


def _redacted_status_payload(args: argparse.Namespace) -> dict[str, Any]:
    settings = _resolve_settings(args)
    active_profile = get_active_profile(settings.config)
    auth = active_profile.get("auth") if isinstance(active_profile, dict) else {}
    stored_key = auth.get("api_key") if isinstance(auth, dict) else None
    return {
        "config_path": str(settings.config_path),
        "active_profile": settings.config.get("active_profile", "default"),
        "api_key": {
            "configured": bool(settings.api_key),
            "redacted": redact_secret(settings.api_key),
            "source": settings.api_key_source,
            "stored": bool(stored_key),
        },
        "base_url": {
            "value": settings.base_url,
            "source": settings.base_url_source,
        },
        "ssh_host": {
            "value": settings.ssh_host,
            "source": settings.ssh_host_source,
        },
    }


def _emit_payload(args: argparse.Namespace, payload: Any, text_lines: Sequence[str]) -> int:
    if getattr(args, "json_output", False):
        _print_json(payload)
    else:
        _print_lines(*text_lines)
    return 0


def _parse_size(value: str) -> int:
    raw = (value or "").strip().lower()
    if not raw:
        raise ValueError("size is required")
    units = (
        ("tb", 1024**4),
        ("gb", 1024**3),
        ("mb", 1024**2),
        ("kb", 1024),
        ("b", 1),
    )
    for suffix, multiplier in units:
        if raw.endswith(suffix) and raw != suffix:
            number = raw[: -len(suffix)].strip()
            return int(float(number) * multiplier)
    return int(raw)


def _human_bytes(value: Optional[int]) -> str:
    amount = float(value or 0)
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if amount < 1024 or unit == "TB":
            if unit == "B":
                return f"{int(amount)}{unit}"
            return f"{amount:.1f}{unit}"
        amount /= 1024
    return f"{int(value or 0)}B"


def _looks_like_uuid(value: str) -> bool:
    try:
        uuid.UUID(str(value))
        return True
    except (TypeError, ValueError, AttributeError):
        return False


def _parse_env_pairs(values: Optional[list[str]]) -> dict[str, str]:
    result: dict[str, str] = {}
    for entry in values or []:
        if "=" not in entry:
            raise ValueError(f"Expected KEY=VALUE, got: {entry}")
        key, value = entry.split("=", 1)
        key = key.strip()
        if not key:
            raise ValueError(f"Invalid env key: {entry}")
        result[key] = value
    return result


def _parse_volume_spec(spec: str) -> dict[str, Any]:
    parts = spec.split(":")
    if len(parts) < 2:
        raise ValueError("volume mounts must use <volume_id>:<mount_path>[:rw|ro[:checkpoint_id]]")
    volume_id = parts[0].strip()
    mount_path = parts[1].strip()
    mode = "rw"
    checkpoint_id = None
    if len(parts) >= 3 and parts[2].strip():
        mode = parts[2].strip().lower()
    if len(parts) >= 4 and parts[3].strip():
        checkpoint_id = parts[3].strip()
    if len(parts) > 4:
        raise ValueError("volume mounts must use <volume_id>:<mount_path>[:rw|ro[:checkpoint_id]]")
    if mode not in {"rw", "ro"}:
        raise ValueError("volume mount mode must be rw or ro")
    if mode == "rw" and checkpoint_id is not None:
        raise ValueError("checkpoint_id is only allowed for ro mounts")
    if mode == "ro" and checkpoint_id is None:
        checkpoint_id = "latest"
    return {
        "volume_id": volume_id,
        "mount_path": mount_path,
        "mode": mode,
        "checkpoint_id": checkpoint_id,
    }


def _find_subparser(parser: argparse.ArgumentParser, path: Sequence[str]) -> Optional[argparse.ArgumentParser]:
    current = parser
    for token in path:
        subparsers_action = next(
            (action for action in current._actions if isinstance(action, argparse._SubParsersAction)),
            None,
        )
        if subparsers_action is None:
            return None
        current = subparsers_action.choices.get(token)
        if current is None:
            return None
    return current


def _add_runtime_options(parser: argparse.ArgumentParser, *, include_json: bool = True, include_ssh_host: bool = False) -> None:
    parser.add_argument("--api-key", help="Override the InstaVM API key for this command")
    parser.add_argument("--base-url", help="Override the InstaVM API base URL for this command")
    if include_ssh_host:
        parser.add_argument("--ssh-host", help="Override the SSH gateway host for this command")
    if include_json:
        parser.add_argument("-j", "--json", dest="json_output", action="store_true", help="Emit JSON output")


def _status_text(payload: dict[str, Any]) -> list[str]:
    lines = [
        f"Config path: {payload['config_path']}",
        f"Active profile: {payload['active_profile']}",
        f"API key: {payload['api_key']['redacted'] or 'not configured'} ({payload['api_key']['source']})",
        f"Base URL: {payload['base_url']['value']} ({payload['base_url']['source']})",
        f"SSH host: {payload['ssh_host']['value']} ({payload['ssh_host']['source']})",
    ]
    if payload["api_key"]["stored"]:
        lines.append("Stored key: present")
    return lines


def handle_help(args: argparse.Namespace, parser: argparse.ArgumentParser) -> int:
    if not args.command_path:
        parser.print_help()
        return 0
    target_parser = _find_subparser(parser, args.command_path)
    if target_parser is None:
        raise ValueError(f"Unknown help topic: {' '.join(args.command_path)}")
    target_parser.print_help()
    return 0


def handle_auth_set_key(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    if sys.stdin.isatty():
        api_key = getpass.getpass("InstaVM API key: ").strip()
    else:
        api_key = sys.stdin.read().strip()
    if not api_key:
        raise ValueError("API key is required. Enter it at the prompt or pipe it over stdin.")

    config_path = update_profile_settings(
        api_key=api_key,
        base_url=args.base_url,
        ssh_host=args.ssh_host,
    )
    payload = {
        "status": "stored",
        "config_path": str(config_path),
        "api_key": redact_secret(api_key),
        "base_url": args.base_url,
        "ssh_host": args.ssh_host,
    }
    text_lines = [
        f"Stored API key in {config_path}",
        f"API key: {payload['api_key']}",
    ]
    if args.base_url:
        text_lines.append(f"Default base URL: {args.base_url}")
    if args.ssh_host:
        text_lines.append(f"Default SSH host: {args.ssh_host}")
    return _emit_payload(args, payload, text_lines)


def handle_auth_status(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    payload = _redacted_status_payload(args)
    return _emit_payload(args, payload, _status_text(payload))


def handle_auth_logout(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    settings = _resolve_settings(args)
    config_path = update_profile_settings(clear_api_key=True)
    payload = {
        "status": "logged_out",
        "config_path": str(config_path),
        "env_key_active": settings.api_key_source == "env",
    }
    text_lines = [f"Removed the stored API key from {config_path}"]
    if payload["env_key_active"]:
        text_lines.append("INSTAVM_API_KEY is still set in the environment and will remain active.")
    return _emit_payload(args, payload, text_lines)


def handle_whoami(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    user = client.get_current_user()
    ssh_keys = client.list_ssh_keys()
    payload = {
        "id": user.get("id"),
        "email": user.get("email"),
        "name": user.get("name"),
        "is_verified": user.get("is_verified"),
        "ssh_keys": ssh_keys,
    }
    text_lines = [
        f"Email: {payload['email'] or '-'}",
        f"User ID: {payload['id']}",
        f"Name: {payload['name'] or '-'}",
        f"Verified: {'yes' if payload['is_verified'] else 'no'}",
        "SSH Keys:",
    ]
    if ssh_keys:
        text_lines.extend(
            f"  {key.get('id')}\t{key.get('fingerprint')}\t{key.get('comment') or ''}".rstrip()
            for key in ssh_keys
        )
    else:
        text_lines.append("  none")
    return _emit_payload(args, payload, text_lines)


def handle_create(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, settings = _require_api_key(args)
    payload: dict[str, Any] = {}
    if args.vm_type == "computer-use":
        payload["image_variant"] = "computer-use"
    if args.timeout is not None:
        payload["vm_lifetime_seconds"] = args.timeout
    if args.memory is not None:
        payload["memory_mb"] = args.memory
    if args.vcpu is not None:
        payload["vcpu_count"] = args.vcpu
    if args.snapshot:
        payload["snapshot_id"] = args.snapshot
    if args.session:
        payload["session_id"] = args.session
    if args.volume:
        payload["volumes"] = [_parse_volume_spec(spec) for spec in args.volume]
    result = client.vms.create(wait=True, **payload)
    text_lines = [
        f"VM: {result.get('vm_id') or '-'}",
        f"Session: {result.get('session_id') or '-'}",
        f"Status: {result.get('status') or '-'}",
    ]
    if result.get("vm_id"):
        text_lines.append(f"SSH: ssh {result['vm_id']}@{settings.ssh_host}")
    return _emit_payload(args, result, text_lines)


def handle_ls(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    vms = client.vms.list_all_records() if args.all else client.vms.list()
    payload = {"vms": vms}
    if not vms:
        return _emit_payload(args, payload, ["No VMs."])
    text_lines = [
        f"{vm.get('vm_id')}\t{vm.get('status') or '-'}\t{vm.get('created_at') or '-'}\t{vm.get('ssh_host') or '-'}"
        for vm in vms
    ]
    return _emit_payload(args, payload, text_lines)


def handle_rm(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    results = []
    for vm_id in args.vm_ids:
        response = client.vms.delete(vm_id)
        results.append({"vm_id": vm_id, **response})
    payload = {"results": results}
    text_lines = [f"{item['vm_id']}: {item.get('status', 'ok')}" for item in results]
    return _emit_payload(args, payload, text_lines)


def handle_clone(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, settings = _require_api_key(args)
    result = client.vms.clone(args.vm_id, wait=True, snapshot_name=args.name)
    text_lines = [
        f"VM: {result.get('vm_id') or '-'}",
        f"Session: {result.get('session_id') or '-'}",
        f"Status: {result.get('status') or '-'}",
    ]
    if result.get("vm_id"):
        text_lines.append(f"SSH: ssh {result['vm_id']}@{settings.ssh_host}")
    return _emit_payload(args, result, text_lines)


def handle_connect(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    settings = _resolve_settings(args)
    ssh_binary = shutil.which("ssh.exe") or shutil.which("ssh")
    if not ssh_binary:
        raise FileNotFoundError("ssh client not found on PATH")
    ssh_args = list(args.ssh_args)
    if ssh_args and ssh_args[0] == "--":
        ssh_args = ssh_args[1:]
    command = [ssh_binary, f"{args.vm_id}@{settings.ssh_host}", *ssh_args]
    if args.json_output:
        _print_json({"command": command})
        return 0
    return subprocess.run(command, check=False).returncode


def handle_snapshot_ls(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    snapshots = client.snapshots.list(snapshot_type=args.snapshot_type)
    payload = {"snapshots": snapshots}
    if not snapshots:
        return _emit_payload(args, payload, ["No snapshots."])
    text_lines = [
        f"{snap.get('id')}\t{snap.get('name') or '-'}\t{snap.get('status') or '-'}\t{snap.get('type') or '-'}"
        for snap in snapshots
    ]
    return _emit_payload(args, payload, text_lines)


def handle_snapshot_create(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    result = client.vms.snapshot(args.vm_id, wait=True, name=args.name)
    text_lines = [
        f"Snapshot: {result.get('snapshot_id') or result.get('id') or '-'}",
        f"Status: {result.get('status') or '-'}",
    ]
    return _emit_payload(args, result, text_lines)


def handle_snapshot_build(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    build_args: dict[str, Any] = {}
    if args.apt:
        build_args["extra_apt_packages"] = args.apt
    if args.pip:
        build_args["extra_pip_packages"] = args.pip
    if args.npm:
        build_args["extra_npm_packages"] = args.npm
    if args.git_clone:
        build_args["git_clone_url"] = args.git_clone
    if args.git_branch:
        build_args["git_clone_branch"] = args.git_branch
    if args.disk_size:
        build_args["disk_size_gb"] = args.disk_size
    if args.run:
        build_args["post_build_cmd"] = args.run
    if args.env:
        build_args["envs"] = _parse_env_pairs(args.env)
    result = client.snapshots.create(
        oci_image=args.oci_image,
        name=args.name,
        vcpu_count=args.vcpu,
        memory_mb=args.memory,
        build_args=build_args or None,
        snapshot_type=args.snapshot_type,
    )
    text_lines = [
        f"Snapshot: {result.get('id') or '-'}",
        f"Name: {result.get('name') or '-'}",
        f"Status: {result.get('status') or '-'}",
    ]
    return _emit_payload(args, result, text_lines)


def handle_snapshot_get(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    result = client.snapshots.get(args.snapshot_id)
    text_lines = [
        f"Snapshot: {result.get('id') or '-'}",
        f"Name: {result.get('name') or '-'}",
        f"Status: {result.get('status') or '-'}",
        f"Type: {result.get('type') or '-'}",
    ]
    return _emit_payload(args, result, text_lines)


def handle_snapshot_rm(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    result = client.snapshots.delete(args.snapshot_id)
    payload = {"snapshot_id": args.snapshot_id, **result}
    text_lines = [f"Deleted snapshot {args.snapshot_id}"]
    return _emit_payload(args, payload, text_lines)


def handle_share_create(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    result = client.shares.create(port=args.port, vm_id=args.vm_id, is_public=args.public)
    text_lines = [
        f"Share: {result.get('share_id')}",
        f"URL: {result.get('url')}",
        f"Port: {result.get('port')}",
        f"Visibility: {'public' if result.get('is_public') else 'private'}",
    ]
    return _emit_payload(args, result, text_lines)


def handle_share_update(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    is_public = args.visibility == "public"
    revoke = args.visibility == "revoke"
    result = client.shares.update(args.share_id, is_public=None if revoke else is_public, revoke=revoke or None)
    payload = {"share_id": args.share_id, **result}
    if revoke:
        text_lines = [f"Revoked share {args.share_id}"]
    else:
        text_lines = [f"Updated share {args.share_id} to {'public' if is_public else 'private'}"]
    return _emit_payload(args, payload, text_lines)


def handle_ssh_key_list(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    keys = client.list_ssh_keys()
    payload = {"keys": keys}
    if not keys:
        return _emit_payload(args, payload, ["No SSH keys."])
    text_lines = [
        f"{key.get('id')}\t{key.get('fingerprint')}\t{key.get('comment') or ''}".rstrip()
        for key in keys
    ]
    return _emit_payload(args, payload, text_lines)


def handle_ssh_key_add(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    public_key = args.public_key
    if public_key is None:
        public_key = sys.stdin.read().strip()
    if not public_key:
        raise ValueError("Public key is required")
    result = client.add_ssh_key(public_key)
    text_lines = [f"Added {result.get('fingerprint') or result.get('id')}"]
    return _emit_payload(args, result, text_lines)


def handle_ssh_key_remove(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    result = client.delete_ssh_key(int(args.key_id))
    payload = {"key_id": int(args.key_id), **result}
    text_lines = [f"Removed {args.key_id}"]
    return _emit_payload(args, payload, text_lines)


def _resolve_desktop_target(client: InstaVM, target: str) -> dict[str, Any]:
    if _looks_like_uuid(target):
        status = client.get_session_status(target)
        status["session_id"] = target
        return status
    vm = client.vms.get(target)
    session_id = vm.get("session_id")
    if session_id:
        status = client.get_session_status(session_id)
        status["session_id"] = session_id
        status.setdefault("vm_id", vm.get("vm_id"))
        return status
    return {
        "session_id": None,
        "vm_id": vm.get("vm_id"),
        "vm_status": vm.get("status"),
        "vm_alive": vm.get("status") == "active",
    }


def _desktop_payload(client: InstaVM, target: str) -> dict[str, Any]:
    payload = _resolve_desktop_target(client, target)
    session_id = payload.get("session_id")
    if payload.get("vm_alive") and session_id:
        try:
            payload.update(client.computer_use.viewer_url(session_id))
        except InstaVMError:
            pass
    return payload


def handle_desktop_status(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    payload = _desktop_payload(client, args.target)
    text_lines = [
        f"Session: {payload.get('session_id') or '-'}",
        f"VM: {payload.get('vm_id') or '-'}",
        f"Status: {payload.get('vm_status') or ('active' if payload.get('vm_alive') else 'stopped')}",
    ]
    if payload.get("viewer_url"):
        text_lines.append(f"Viewer: {payload['viewer_url']}")
    return _emit_payload(args, payload, text_lines)


def handle_desktop_viewer(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    payload = _desktop_payload(client, args.target)
    if not payload.get("session_id") or not payload.get("vm_alive"):
        raise ValueError("Desktop is not running.")
    if "viewer_url" not in payload:
        payload.update(client.computer_use.viewer_url(payload["session_id"]))
    text_lines = [
        f"Session: {payload.get('session_id')}",
        f"VM: {payload.get('vm_id') or '-'}",
        f"Viewer URL: {payload.get('viewer_url')}",
    ]
    if payload.get("websocket_url"):
        text_lines.append(f"WebSocket: {payload['websocket_url']}")
    return _emit_payload(args, payload, text_lines)


def handle_desktop_start(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    payload = _desktop_payload(client, args.target)
    if payload.get("vm_alive"):
        return _emit_payload(
            args,
            payload,
            [
                f"Desktop already running: {payload.get('vm_id') or '-'}",
                f"Session: {payload.get('session_id') or '-'}",
            ],
        )
    session_id = payload.get("session_id")
    if not session_id:
        raise ValueError("Desktop start requires a session-backed VM or session ID.")
    result = client.vms.create(wait=True, session_id=session_id, image_variant="computer-use")
    merged = _desktop_payload(client, result.get("session_id") or session_id)
    merged.update(result)
    text_lines = [
        f"Desktop started: {merged.get('vm_id') or result.get('vm_id') or '-'}",
        f"Session: {merged.get('session_id') or result.get('session_id') or '-'}",
    ]
    if merged.get("viewer_url"):
        text_lines.append(f"Viewer: {merged['viewer_url']}")
    return _emit_payload(args, merged, text_lines)


def handle_desktop_stop(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    payload = _desktop_payload(client, args.target)
    if not payload.get("vm_alive"):
        return _emit_payload(args, payload, ["Desktop is already stopped."])
    if payload.get("session_id"):
        result = client.kill(payload["session_id"])
    elif payload.get("vm_id"):
        result = client.vms.delete(payload["vm_id"])
    else:
        raise ValueError("No active desktop VM found.")
    merged = {"session_id": payload.get("session_id"), "vm_id": payload.get("vm_id"), **result}
    text_lines = [f"Stopping {payload.get('vm_id') or payload.get('session_id')}..."]
    return _emit_payload(args, merged, text_lines)


def handle_volume_ls(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    volumes = client.volumes.list(refresh_usage=args.refresh)
    payload = {"volumes": volumes}
    if not volumes:
        return _emit_payload(args, payload, ["No volumes."])
    text_lines = [
        f"{volume.get('id')}\t{volume.get('name')}\t{_human_bytes(volume.get('used_bytes'))}/{_human_bytes(volume.get('quota_bytes'))}\t{volume.get('status')}"
        for volume in volumes
    ]
    return _emit_payload(args, payload, text_lines)


def handle_volume_get(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    volume = client.volumes.get(args.volume_id, refresh_usage=args.refresh)
    text_lines = [
        f"ID: {volume.get('id')}",
        f"Name: {volume.get('name')}",
        f"Status: {volume.get('status')}",
        f"Quota: {_human_bytes(volume.get('quota_bytes'))}",
        f"Used: {_human_bytes(volume.get('used_bytes'))}",
    ]
    return _emit_payload(args, volume, text_lines)


def handle_volume_create(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    result = client.volumes.create(name=args.name, quota_bytes=_parse_size(args.quota))
    text_lines = [
        f"Volume: {result.get('id')}",
        f"Name: {result.get('name')}",
        f"Quota: {_human_bytes(result.get('quota_bytes'))}",
    ]
    return _emit_payload(args, result, text_lines)


def handle_volume_update(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    kwargs: dict[str, Any] = {}
    if args.name:
        kwargs["name"] = args.name
    if args.quota:
        kwargs["quota_bytes"] = _parse_size(args.quota)
    if not kwargs:
        raise ValueError("Provide --name and/or --quota")
    result = client.volumes.update(args.volume_id, **kwargs)
    text_lines = [f"Updated volume {result.get('id') or args.volume_id}"]
    return _emit_payload(args, result, text_lines)


def handle_volume_rm(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    result = client.volumes.delete(args.volume_id)
    payload = {"volume_id": args.volume_id, **result}
    return _emit_payload(args, payload, [f"Deleted {args.volume_id}"])


def handle_volume_checkpoint_ls(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    checkpoints = client.volumes.list_checkpoints(args.volume_id)
    payload = {"checkpoints": checkpoints}
    if not checkpoints:
        return _emit_payload(args, payload, ["No checkpoints."])
    text_lines = [
        f"{checkpoint.get('id')}\t{checkpoint.get('name') or '-'}\t{checkpoint.get('status') or '-'}"
        for checkpoint in checkpoints
    ]
    return _emit_payload(args, payload, text_lines)


def handle_volume_checkpoint_create(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    result = client.volumes.create_checkpoint(args.volume_id, name=args.name)
    return _emit_payload(args, result, [f"Created checkpoint {result.get('id') or '-'}"])


def handle_volume_checkpoint_rm(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    result = client.volumes.delete_checkpoint(args.volume_id, args.checkpoint_id)
    payload = {"checkpoint_id": args.checkpoint_id, **result}
    return _emit_payload(args, payload, [f"Deleted checkpoint {args.checkpoint_id}"])


def handle_volume_files_ls(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    files = client.volumes.list_files(
        args.volume_id,
        prefix=args.prefix or "",
        recursive=not args.no_recursive,
        limit=args.limit,
    )
    payload = {"files": files}
    if not files:
        return _emit_payload(args, payload, ["No files."])
    text_lines = [f"{entry.get('path')}\t{_human_bytes(entry.get('size'))}" for entry in files]
    return _emit_payload(args, payload, text_lines)


def handle_volume_files_upload(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    remote_path = args.path or Path(args.local_path).name
    result = client.volumes.upload_file(
        args.volume_id,
        file_path=args.local_path,
        path=remote_path,
        overwrite=args.overwrite,
    )
    return _emit_payload(args, result, [f"Uploaded {args.local_path} -> {remote_path}"])


def handle_volume_files_download(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    output_path = args.out or Path(args.remote_path).name
    result = client.volumes.download_file(args.volume_id, path=args.remote_path, local_path=output_path)
    payload = {
        "path": result.get("path"),
        "filename": result.get("filename"),
        "size": result.get("size") or result.get("saved_bytes"),
        "local_path": result.get("local_path") or output_path,
    }
    return _emit_payload(args, payload, [f"Downloaded {args.remote_path} -> {payload['local_path']}"])


def handle_volume_files_rm(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    client, _settings = _require_api_key(args)
    result = client.volumes.delete_file(args.volume_id, path=args.remote_path)
    payload = {"path": args.remote_path, **result}
    return _emit_payload(args, payload, [f"Deleted {args.remote_path}"])


def handle_doc(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    payload = {"url": DEFAULT_DOCS_URL}
    return _emit_payload(args, payload, [DEFAULT_DOCS_URL])


def handle_billing(args: argparse.Namespace, _parser: argparse.ArgumentParser) -> int:
    payload = {"url": DEFAULT_BILLING_URL}
    return _emit_payload(args, payload, [DEFAULT_BILLING_URL])


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="instavm",
        description="InstaVM CLI for VM, snapshot, volume, desktop, identity, and sharing workflows.",
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    help_parser = subparsers.add_parser("help", help="Show help for the CLI or a subcommand")
    help_parser.add_argument("command_path", nargs="*", help="Command path, for example: volume files")
    help_parser.set_defaults(handler=handle_help)

    auth_parser = subparsers.add_parser("auth", help="Manage stored auth for the installed CLI")
    auth_subparsers = auth_parser.add_subparsers(dest="auth_command", required=True)

    auth_set_key = auth_subparsers.add_parser(
        "set-key",
        help="Store an InstaVM API key in ~/.instavm/config.json",
        description="Store an InstaVM API key from a hidden prompt or piped stdin.",
    )
    auth_set_key.add_argument("--base-url", help="Persist a default API base URL")
    auth_set_key.add_argument("--ssh-host", help="Persist a default SSH host")
    auth_set_key.add_argument("-j", "--json", dest="json_output", action="store_true", help="Emit JSON output")
    auth_set_key.set_defaults(handler=handle_auth_set_key)

    auth_status = auth_subparsers.add_parser("status", help="Show the effective CLI auth and endpoint settings")
    _add_runtime_options(auth_status, include_json=True, include_ssh_host=True)
    auth_status.set_defaults(handler=handle_auth_status)

    auth_logout = auth_subparsers.add_parser("logout", help="Remove the stored API key from the local CLI config")
    _add_runtime_options(auth_logout, include_json=True, include_ssh_host=True)
    auth_logout.set_defaults(handler=handle_auth_logout)

    whoami = subparsers.add_parser("whoami", help="Show the current account and SSH key registrations")
    _add_runtime_options(whoami)
    whoami.set_defaults(handler=handle_whoami)

    create = subparsers.add_parser("create", aliases=["new"], help="Create a new VM or desktop")
    _add_runtime_options(create)
    create.add_argument("--type", dest="vm_type", choices=["standard", "computer-use"], default="standard")
    create.add_argument("--timeout", type=int, help="VM lifetime in seconds")
    create.add_argument("--memory", type=int, help="Memory in MB")
    create.add_argument("--vcpu", type=int, help="vCPU count")
    create.add_argument("--snapshot", help="Snapshot ID to boot from")
    create.add_argument("--session", help="Session ID to attach or respawn")
    create.add_argument("--volume", action="append", help="Volume spec: <volume_id>:<mount_path>[:rw|ro[:checkpoint_id]]")
    create.set_defaults(handler=handle_create)

    ls_parser = subparsers.add_parser("ls", aliases=["list"], help="List your VMs")
    _add_runtime_options(ls_parser)
    ls_parser.add_argument("--all", "-a", action="store_true", help="Include all VM records, not just active VMs")
    ls_parser.set_defaults(handler=handle_ls)

    rm_parser = subparsers.add_parser("rm", aliases=["delete"], help="Delete one or more VMs")
    _add_runtime_options(rm_parser)
    rm_parser.add_argument("vm_ids", nargs="+", help="VM IDs to terminate")
    rm_parser.set_defaults(handler=handle_rm)

    clone_parser = subparsers.add_parser("clone", help="Clone a VM via snapshot")
    _add_runtime_options(clone_parser)
    clone_parser.add_argument("vm_id", help="Source VM ID")
    clone_parser.add_argument("--name", help="Snapshot name used during clone")
    clone_parser.set_defaults(handler=handle_clone)

    connect_parser = subparsers.add_parser("connect", help="SSH into a VM with the local ssh client")
    _add_runtime_options(connect_parser, include_json=True, include_ssh_host=True)
    connect_parser.add_argument("vm_id", help="VM ID to connect to")
    connect_parser.add_argument("ssh_args", nargs=argparse.REMAINDER, help="Additional ssh arguments after --")
    connect_parser.set_defaults(handler=handle_connect)

    snapshot_parser = subparsers.add_parser("snapshot", help="List, inspect, build, create, and delete snapshots")
    snapshot_subparsers = snapshot_parser.add_subparsers(dest="snapshot_command", required=True)

    snapshot_ls = snapshot_subparsers.add_parser("ls", aliases=["list"], help="List snapshots")
    _add_runtime_options(snapshot_ls)
    snapshot_ls.add_argument("--type", dest="snapshot_type", choices=["user", "system"], help="Snapshot type filter")
    snapshot_ls.set_defaults(handler=handle_snapshot_ls)

    snapshot_create = snapshot_subparsers.add_parser("create", help="Create a snapshot from a running VM")
    _add_runtime_options(snapshot_create)
    snapshot_create.add_argument("vm_id", help="VM ID to snapshot")
    snapshot_create.add_argument("--name", help="Snapshot name")
    snapshot_create.set_defaults(handler=handle_snapshot_create)

    snapshot_build = snapshot_subparsers.add_parser("build", help="Build a snapshot from an OCI image")
    _add_runtime_options(snapshot_build)
    snapshot_build.add_argument("oci_image", help="OCI image reference")
    snapshot_build.add_argument("--name", help="Snapshot name")
    snapshot_build.add_argument("--vcpu", type=int, help="vCPU count")
    snapshot_build.add_argument("--memory", type=int, help="Memory in MB")
    snapshot_build.add_argument("--disk-size", type=int, help="Snapshot disk size in GB")
    snapshot_build.add_argument("--env", action="append", help="Build environment variable in KEY=VALUE form")
    snapshot_build.add_argument("--git-clone", help="Git repository to clone during build")
    snapshot_build.add_argument("--git-branch", help="Git branch to checkout")
    snapshot_build.add_argument("--run", help="Command to run after the image is prepared")
    snapshot_build.add_argument("--apt", help="Extra apt packages")
    snapshot_build.add_argument("--pip", help="Extra pip packages")
    snapshot_build.add_argument("--npm", help="Extra npm packages")
    snapshot_build.add_argument("--type", dest="snapshot_type", choices=["user", "system"], default="user")
    snapshot_build.set_defaults(handler=handle_snapshot_build)

    snapshot_get = snapshot_subparsers.add_parser("get", help="Get a snapshot by ID")
    _add_runtime_options(snapshot_get)
    snapshot_get.add_argument("snapshot_id", help="Snapshot ID")
    snapshot_get.set_defaults(handler=handle_snapshot_get)

    snapshot_rm = snapshot_subparsers.add_parser("rm", aliases=["delete"], help="Delete a snapshot")
    _add_runtime_options(snapshot_rm)
    snapshot_rm.add_argument("snapshot_id", help="Snapshot ID")
    snapshot_rm.set_defaults(handler=handle_snapshot_rm)

    share_parser = subparsers.add_parser("share", help="Create and manage VM shares")
    share_subparsers = share_parser.add_subparsers(dest="share_command", required=True)

    share_create = share_subparsers.add_parser("create", help="Create a share for a VM and port")
    _add_runtime_options(share_create)
    share_create.add_argument("vm_id", help="VM ID")
    share_create.add_argument("port", type=int, help="Port to expose")
    share_create.add_argument("--public", action="store_true", help="Create the share as public")
    share_create.set_defaults(handler=handle_share_create)

    share_set_public = share_subparsers.add_parser("set-public", help="Set a share to public visibility")
    _add_runtime_options(share_set_public)
    share_set_public.add_argument("share_id", help="Share ID")
    share_set_public.set_defaults(handler=handle_share_update, visibility="public")

    share_set_private = share_subparsers.add_parser("set-private", help="Set a share to private visibility")
    _add_runtime_options(share_set_private)
    share_set_private.add_argument("share_id", help="Share ID")
    share_set_private.set_defaults(handler=handle_share_update, visibility="private")

    share_revoke = share_subparsers.add_parser("revoke", help="Revoke a share by share ID")
    _add_runtime_options(share_revoke)
    share_revoke.add_argument("share_id", help="Share ID")
    share_revoke.set_defaults(handler=handle_share_update, visibility="revoke")

    ssh_key_parser = subparsers.add_parser("ssh-key", aliases=["sshkey"], help="Manage SSH keys on your account")
    ssh_key_subparsers = ssh_key_parser.add_subparsers(dest="ssh_key_command", required=True)

    ssh_key_list = ssh_key_subparsers.add_parser("list", help="List active SSH keys")
    _add_runtime_options(ssh_key_list)
    ssh_key_list.set_defaults(handler=handle_ssh_key_list)

    ssh_key_add = ssh_key_subparsers.add_parser("add", help="Add an SSH public key")
    _add_runtime_options(ssh_key_add)
    ssh_key_add.add_argument("public_key", nargs="?", help="Public key string. If omitted, reads from stdin.")
    ssh_key_add.set_defaults(handler=handle_ssh_key_add)

    ssh_key_remove = ssh_key_subparsers.add_parser("remove", help="Remove an SSH key by key ID")
    _add_runtime_options(ssh_key_remove)
    ssh_key_remove.add_argument("key_id", help="SSH key ID")
    ssh_key_remove.set_defaults(handler=handle_ssh_key_remove)

    desktop_parser = subparsers.add_parser("desktop", help="Manage computer-use desktops and viewer handoff")
    desktop_subparsers = desktop_parser.add_subparsers(dest="desktop_command", required=True)
    for name, help_text, handler in (
        ("status", "Show desktop status for a VM or session", handle_desktop_status),
        ("start", "Start a computer-use desktop for a session", handle_desktop_start),
        ("stop", "Stop a running desktop VM", handle_desktop_stop),
        ("viewer", "Get the noVNC viewer URL for a running desktop", handle_desktop_viewer),
    ):
        desktop_command = desktop_subparsers.add_parser(name, help=help_text)
        _add_runtime_options(desktop_command)
        desktop_command.add_argument("target", help="VM ID or session ID")
        desktop_command.set_defaults(handler=handler)

    volume_parser = subparsers.add_parser("volume", aliases=["volumes"], help="Manage persistent volumes")
    volume_subparsers = volume_parser.add_subparsers(dest="volume_command", required=True)

    volume_ls = volume_subparsers.add_parser("ls", aliases=["list"], help="List volumes")
    _add_runtime_options(volume_ls)
    volume_ls.add_argument("--refresh", action="store_true", help="Refresh usage before listing")
    volume_ls.set_defaults(handler=handle_volume_ls)

    volume_get = volume_subparsers.add_parser("get", help="Get a volume")
    _add_runtime_options(volume_get)
    volume_get.add_argument("volume_id", help="Volume ID")
    volume_get.add_argument("--refresh", action="store_true", help="Refresh usage before reading")
    volume_get.set_defaults(handler=handle_volume_get)

    volume_create = volume_subparsers.add_parser("create", help="Create a volume")
    _add_runtime_options(volume_create)
    volume_create.add_argument("name", help="Volume name")
    volume_create.add_argument("--quota", required=True, help="Quota in bytes or units like 5gb")
    volume_create.set_defaults(handler=handle_volume_create)

    volume_update = volume_subparsers.add_parser("update", help="Update a volume")
    _add_runtime_options(volume_update)
    volume_update.add_argument("volume_id", help="Volume ID")
    volume_update.add_argument("--name", help="Updated volume name")
    volume_update.add_argument("--quota", help="Updated quota in bytes or units like 5gb")
    volume_update.set_defaults(handler=handle_volume_update)

    volume_rm = volume_subparsers.add_parser("rm", aliases=["delete"], help="Delete a volume")
    _add_runtime_options(volume_rm)
    volume_rm.add_argument("volume_id", help="Volume ID")
    volume_rm.set_defaults(handler=handle_volume_rm)

    checkpoint_parser = volume_subparsers.add_parser("checkpoint", help="Manage checkpoints for a volume")
    checkpoint_subparsers = checkpoint_parser.add_subparsers(dest="volume_checkpoint_command", required=True)

    checkpoint_ls = checkpoint_subparsers.add_parser("ls", aliases=["list"], help="List checkpoints for a volume")
    _add_runtime_options(checkpoint_ls)
    checkpoint_ls.add_argument("volume_id", help="Volume ID")
    checkpoint_ls.set_defaults(handler=handle_volume_checkpoint_ls)

    checkpoint_create = checkpoint_subparsers.add_parser("create", help="Create a checkpoint for a volume")
    _add_runtime_options(checkpoint_create)
    checkpoint_create.add_argument("volume_id", help="Volume ID")
    checkpoint_create.add_argument("--name", help="Checkpoint name")
    checkpoint_create.set_defaults(handler=handle_volume_checkpoint_create)

    checkpoint_rm = checkpoint_subparsers.add_parser("rm", aliases=["delete"], help="Delete a checkpoint")
    _add_runtime_options(checkpoint_rm)
    checkpoint_rm.add_argument("volume_id", help="Volume ID")
    checkpoint_rm.add_argument("checkpoint_id", help="Checkpoint ID")
    checkpoint_rm.set_defaults(handler=handle_volume_checkpoint_rm)

    files_parser = volume_subparsers.add_parser("files", help="Manage files stored inside a volume")
    files_subparsers = files_parser.add_subparsers(dest="volume_files_command", required=True)

    files_ls = files_subparsers.add_parser("ls", aliases=["list"], help="List files inside a volume")
    _add_runtime_options(files_ls)
    files_ls.add_argument("volume_id", help="Volume ID")
    files_ls.add_argument("--prefix", help="Path prefix filter")
    files_ls.add_argument("--limit", type=int, default=1000, help="Maximum number of entries")
    files_ls.add_argument("--no-recursive", action="store_true", help="Do not recurse into subdirectories")
    files_ls.set_defaults(handler=handle_volume_files_ls)

    files_upload = files_subparsers.add_parser("upload", help="Upload a file into a volume")
    _add_runtime_options(files_upload)
    files_upload.add_argument("volume_id", help="Volume ID")
    files_upload.add_argument("local_path", help="Local file path")
    files_upload.add_argument("--path", help="Remote path inside the volume")
    files_upload.add_argument("--overwrite", action="store_true", help="Overwrite if the path already exists")
    files_upload.set_defaults(handler=handle_volume_files_upload)

    files_download = files_subparsers.add_parser("download", help="Download a file from a volume")
    _add_runtime_options(files_download)
    files_download.add_argument("volume_id", help="Volume ID")
    files_download.add_argument("remote_path", help="Remote path inside the volume")
    files_download.add_argument("--out", help="Local output path")
    files_download.set_defaults(handler=handle_volume_files_download)

    files_rm = files_subparsers.add_parser("rm", aliases=["delete"], help="Delete a file from a volume")
    _add_runtime_options(files_rm)
    files_rm.add_argument("volume_id", help="Volume ID")
    files_rm.add_argument("remote_path", help="Remote path inside the volume")
    files_rm.set_defaults(handler=handle_volume_files_rm)

    doc_parser = subparsers.add_parser("doc", aliases=["docs"], help="Show InstaVM documentation links")
    doc_parser.add_argument("-j", "--json", dest="json_output", action="store_true", help="Emit JSON output")
    doc_parser.set_defaults(handler=handle_doc)

    billing_parser = subparsers.add_parser("billing", help="Show the billing portal URL")
    billing_parser.add_argument("-j", "--json", dest="json_output", action="store_true", help="Emit JSON output")
    billing_parser.set_defaults(handler=handle_billing)

    return parser


def main(argv: Optional[Sequence[str]] = None) -> int:
    parser = build_parser()
    try:
        args = parser.parse_args(argv)
        return args.handler(args, parser)
    except (AuthenticationError, ConfigError, FileNotFoundError, InstaVMError, ValueError) as exc:
        sys.stderr.write(f"Error: {exc}\n")
        return 1
    except KeyboardInterrupt:
        sys.stderr.write("Interrupted.\n")
        return 130


if __name__ == "__main__":
    raise SystemExit(main())
