#!/usr/bin/env python3
# Copyright (C) 2024 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the 'License');
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an 'AS IS' BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
import subprocess
import sys


def list_instances(project_id):
  try:
    result = subprocess.run([
        'gcloud', 'compute', 'instances', 'list', '--project', project_id,
        '--format', 'table(name,zone)'
    ],
                            check=True,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE,
                            text=True)
    lines = result.stdout.strip().split('\n')
    instances = [tuple(line.split()) for line in lines[1:]]  # Skip the header
    return instances
  except subprocess.CalledProcessError as e:
    print(f'Error retrieving instances: {e.stderr}')
    sys.exit(1)


def main():
  DEFAULT_PROJECT_ID = 'perfetto-ci'
  # project_id = os.getenv('CLOUDSDK_CORE_PROJECT', DEFAULT_PROJECT_ID)

  parser = argparse.ArgumentParser()
  parser.add_argument(
      '-p',
      '--project-id',
      metavar='PROJECT_ID',
      required=False,
      help='The Cloud project id. Defaults to CLOUDSDK_CORE_PROJECT',
      default=os.getenv('CLOUDSDK_CORE_PROJECT', DEFAULT_PROJECT_ID))
  args = parser.parse_args()
  project_id = args.project_id

  print('Using Cloud project: %s' % project_id)
  print('If this script fail ensure that:')
  print(' - The cloud project has been configured as per go/gce-beyondcorp-ssh')
  print(' - Register your key as per "Ensure that you are registered with OS')

  instances = list_instances(project_id)
  if not instances:
    print('No GCE instances found.')
    sys.exit(0)

  print('Available VMs:')
  for idx, (name, zone) in enumerate(instances, start=1):
    print(f'{idx}. {name} ({zone})')

  try:
    vm_number = int(input('Enter the number of the VM you want to ssh into: '))
    if vm_number < 1 or vm_number > len(instances):
      raise ValueError
  except ValueError:
    print('Invalid selection. Please run the script again.')
    sys.exit(1)

  # Get the selected VM's name and zone
  selected_instance = instances[vm_number - 1]
  vm_name, vm_zone = selected_instance
  user = os.getenv('USER', 'username')
  ssh_arg = '%s_google_com@nic0.%s.%s.c.%s.internal.gcpnode.com' % (
      user, vm_name, vm_zone, project_id)
  print('ssh ' + ssh_arg)
  os.execvp('ssh', ['ssh', ssh_arg])


if __name__ == '__main__':
  main()