# Copyright 2018 Google Inc. All Rights Reserved.
#
# 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.
"""Utilities shared by cloud-shell commands."""

from __future__ import absolute_import
from __future__ import unicode_literals
import base64

from googlecloudsdk.api_lib.util import apis
from googlecloudsdk.api_lib.util import waiter
from googlecloudsdk.command_lib.util.ssh import ssh
from googlecloudsdk.core import exceptions
from googlecloudsdk.core import log

DEFAULT_ENVIRONMENT_NAME = 'users/me/environments/default'


class UnsupportedPlatform(exceptions.Error):
  """Raised when attempting to run a command on an unsupported platform."""


def ParseCommonArgs(parser):
  """Parses arguments common to all cloud-shell commands."""

  parser.add_argument(
      '--force-key-file-overwrite',
      help="""\
      If enabled gcloud will regenerate and overwrite the files associated
      with a broken SSH key without asking for confirmation in both
      interactive and non-interactive environment.

      If disabled gcloud will not attempt to regenerate the files associated
      with a broken SSH key and fail in both interactive and non-interactive
      environment.
      """,
      action='store_true')
  parser.add_argument(
      '--ssh-key-file',
      help="""\
      The path to the SSH key file. By default, this is
        *~/.ssh/google_compute_engine*.
      """,
      action='store_true')


def PrepareEnvironment(args):
  """Ensures that the user's environment is ready to accept SSH connections."""

  # Load Cloud Shell API
  client = apis.GetClientInstance('cloudshell', 'v1alpha1')
  messages = apis.GetMessagesModule('cloudshell', 'v1alpha1')
  operations_client = apis.GetClientInstance('cloudshell', 'v1')

  # Ensure we have a key pair on the local machine
  ssh_env = ssh.Environment.Current()
  ssh_env.RequireSSH()
  keys = ssh.Keys.FromFilename(filename=args.ssh_key_file)
  keys.EnsureKeysExist(overwrite=args.force_key_file_overwrite)

  # Look up the Cloud Shell environment
  environment = client.users_environments.Get(
      messages.CloudshellUsersEnvironmentsGetRequest(
          name=DEFAULT_ENVIRONMENT_NAME))

  # If the environment doesn't have the public key, push it
  key_parts = keys.GetPublicKey().ToEntry().split(' ')
  key = messages.PublicKey(
      format=messages.PublicKey.FormatValueValuesEnum(key_parts[0].replace(
          '-', '_').upper()),
      key=base64.b64decode(key_parts[1]),
  )
  has_key = False
  for candidate in environment.publicKeys:
    if key.format == candidate.format and key.key == candidate.key:
      has_key = True
      break
  if not has_key:
    log.Print('Pushing your public key to Cloud Shell...')
    client.users_environments_publicKeys.Create(
        messages.CloudshellUsersEnvironmentsPublicKeysCreateRequest(
            parent=DEFAULT_ENVIRONMENT_NAME,
            createPublicKeyRequest=messages.CreatePublicKeyRequest(key=key),
        ))

  # If the environment isn't running, start it
  if environment.state != messages.Environment.StateValueValuesEnum.RUNNING:
    log.Print('Starting your Cloud Shell machine...')
    start_operation = client.users_environments.Start(
        messages.CloudshellUsersEnvironmentsStartRequest(
            name=DEFAULT_ENVIRONMENT_NAME))

    environment = waiter.WaitFor(
        StartEnvironmentPoller(client.users_environments,
                               operations_client.operations),
        start_operation,
        'Waiting for your Cloud Shell machine to start',
        sleep_ms=500,
        max_wait_ms=None)

  return ConnectionInfo(
      ssh_env=ssh_env,
      user=environment.sshUsername,
      host=environment.sshHost,
      port=environment.sshPort,
      key=keys.key_file,
  )


class ConnectionInfo(object):

  def __init__(self, ssh_env, user, host, port, key):
    self.ssh_env = ssh_env
    self.user = user
    self.host = host
    self.port = port
    self.key = key


class StartEnvironmentPoller(waiter.OperationPoller):
  """Poller for start environment operations."""

  def __init__(self, environments_service, operations_service):
    self.environments_service = environments_service
    self.operations_service = operations_service

  def IsDone(self, operation):
    return operation.done

  def Poll(self, operation):
    request_type = self.operations_service.GetRequestType('Get')
    return self.operations_service.Get(request_type(name=operation.name))

  def GetResult(self, operation):
    request_type = self.environments_service.GetRequestType('Get')
    return self.environments_service.Get(
        request_type(name=DEFAULT_ENVIRONMENT_NAME))
