# -*- coding: utf-8 -*- #
# Copyright 2018 Google LLC. 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."""


import argparse
import datetime
import time

from apitools.base.py import exceptions as apitools_exceptions
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 as core_exceptions
from googlecloudsdk.core import log
from googlecloudsdk.core.credentials import store


DEFAULT_ENVIRONMENT_NAME = 'users/me/environments/default'

MIN_CREDS_EXPIRY = datetime.timedelta(minutes=30)
MIN_CREDS_EXPIRY_SECONDS = '{}s'.format(MIN_CREDS_EXPIRY.seconds)


class UnsupportedPlatform(core_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 AddSshArgFlag(parser):
  parser.add_argument(
      'ssh_args',
      nargs=argparse.REMAINDER,
      help="""\
          Flags and positionals passed to the underlying ssh implementation.
          """,
      example="""\
        $ {command} -- -vvv
      """)


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

  # Load Cloud Shell API.
  client = apis.GetClientInstance('cloudshell', 'v1')
  messages = apis.GetMessagesModule('cloudshell', 'v1')
  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 = keys.GetPublicKey().ToEntry()
  has_key = False
  for candidate in environment.publicKeys:
    if key == candidate:
      has_key = True
      break
  add_public_key_error = None
  if not has_key:
    try:
      add_public_key_operation = client.users_environments.AddPublicKey(
          messages.CloudshellUsersEnvironmentsAddPublicKeyRequest(
              environment=DEFAULT_ENVIRONMENT_NAME,
              addPublicKeyRequest=messages.AddPublicKeyRequest(key=key),
          )
      )

      environment = waiter.WaitFor(
          EnvironmentPoller(
              client.users_environments, operations_client.operations
          ),
          add_public_key_operation,
          'Pushing your public key to Cloud Shell',
          sleep_ms=500,
          max_wait_ms=None,
      )
      # Wait for the public key to propagate. Unfortunately there is no
      # api we can check for this.
      log.Print('Waiting for your public key to propagate...')
      time.sleep(5)
      key_file = keys.key_file
    except apitools_exceptions.HttpError as e:
      if e.status_code == 410:
        log.Print('AddPublicKey is deprecated. Skipping public key push.')
        key_file = None
        add_public_key_error = e
      else:
        raise
  else:
    key_file = keys.key_file

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

    access_token = None
    if args.IsKnownAndSpecified('authorize_session') and args.authorize_session:
      access_token = store.GetFreshAccessTokenIfEnabled(
          min_expiry_duration=MIN_CREDS_EXPIRY_SECONDS
      )

    start_operation = client.users_environments.Start(
        messages.CloudshellUsersEnvironmentsStartRequest(
            name=DEFAULT_ENVIRONMENT_NAME,
            startEnvironmentRequest=messages.StartEnvironmentRequest(
                accessToken=access_token
            ),
        )
    )

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

  if not environment.sshHost:
    _RaiseExceptionDueToStartOperationError(start_operation)

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


def _GetDetail(details, type_class):
  if details is None:
    return None

  for detail in details:
    for additional_property in detail.additionalProperties:
      if (
          additional_property.key == '@type'
          and additional_property.value.string_value == type_class
      ):
        return detail
  return None


def _GetAdditionalPropertyFromDetail(detail, key):
  if detail is None:
    return None

  for additional_property in detail.additionalProperties:
    if additional_property.key == key:
      return additional_property.value
  return None


def _RaiseExceptionDueToStartOperationError(operation):
  """Raises a core_exceptions.Error based on the start operation's error details.

  If the error is due to an unverified account, a detailed error including the
  verification URL (if it exists) will be raised. Otherwise, a generic error is
  raised.

  Args:
    operation: The operation object returned from the Cloud Shell API.
  """
  details = None
  if operation and operation.error:
    details = operation.error.details

  error_info = _GetDetail(details, 'type.googleapis.com/google.rpc.ErrorInfo')
  reason_property = _GetAdditionalPropertyFromDetail(error_info, 'reason')
  error_reason = reason_property.string_value if reason_property else None

  if error_reason != 'ACCOUNT_UNVERIFIED':
    raise core_exceptions.Error('The Cloud Shell machine did not start.')

  verify_url = None
  help_message = _GetDetail(details, 'type.googleapis.com/google.rpc.Help')

  # Attempt to retrieve the uplevel link from the operation help message.
  # If it doesn't exist, use the Cloud Shell homepage as a fallback.
  if help_message:
    links = _GetAdditionalPropertyFromDetail(help_message, 'links')
    if links:
      for entries in links.array_value.entries:
        if verify_url:
          break
        for link in entries.object_value.properties:
          if link.key == 'url':
            verify_url = link.value.string_value
            break
  if verify_url:
    raise core_exceptions.Error(
        'Your account is unverified. Please verify your account at'
        f' {verify_url} or at https://shell.cloud.google.com.'
    )
  else:
    raise core_exceptions.Error(
        'Your account is unverified. Please verify your account at'
        ' https://shell.cloud.google.com.'
    )


def AuthorizeEnvironment():
  """Pushes gcloud command-line tool credentials to the user's environment."""

  client = apis.GetClientInstance('cloudshell', 'v1')
  messages = apis.GetMessagesModule('cloudshell', 'v1')

  # Load creds and refresh them if they're about to expire.
  access_token = store.GetFreshAccessTokenIfEnabled(
      min_expiry_duration=MIN_CREDS_EXPIRY_SECONDS)
  if access_token:
    client.users_environments.Authorize(
        messages.CloudshellUsersEnvironmentsAuthorizeRequest(
            name=DEFAULT_ENVIRONMENT_NAME,
            authorizeEnvironmentRequest=messages.AuthorizeEnvironmentRequest(
                accessToken=access_token)))


def GenerateAccessToken(environment=DEFAULT_ENVIRONMENT_NAME):
  """Generates an access token for the user's environment."""
  client = apis.GetClientInstance('cloudshell', 'v1')
  messages = apis.GetMessagesModule('cloudshell', 'v1')
  req = messages.CloudshellUsersEnvironmentsGenerateAccessTokenRequest(
      environment=environment
  )
  res = client.users_environments.GenerateAccessToken(req)
  return res.accessToken


class ConnectionInfo(object):
  """Connection information for Cloud Shell."""

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


class EnvironmentPoller(waiter.OperationPoller):
  """Poller for 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))