#!/usr/bin/env python
"""The classes used to define config used to delegate BQ commands to gcloud."""

from collections.abc import Callable
import sys
from typing import Dict, List, Optional, Union
from typing_extensions import TypeAlias

PrimitiveFlagValue: TypeAlias = Union[str, bool, int]
NestedStrDict: TypeAlias = Dict[str, Union[str, 'NestedStrDict']]


def _flatten_flag_dictionary(
    mapped_flags: Dict[str, PrimitiveFlagValue],
) -> List[str]:
  """Returns the gcloud command flags as an array of strings."""
  flag_array: List[str] = []
  for name, value in mapped_flags.items():
    if not isinstance(value, bool):
      flag_array.append(f'--{name}={str(value)}')
    elif value:
      flag_array.append(f'--{name}')
    else:
      flag_array.append(f'--no-{name}')
  return flag_array


def quote_flag_values(command_array: List[str]) -> List[str]:
  """Returns the gcloud command flags after quoting the flag values."""
  result: List[str] = []
  for command_or_flag in command_array:
    if command_or_flag.startswith('--') and '=' in command_or_flag:
      (name, _, value) = command_or_flag.partition('=')
      result.append(f"{name}='{value}'")
    else:
      result.append(command_or_flag)
  return result


class BigqueryGcloudDelegationUserError(Exception):
  """Class to represent a user error during gcloud delegation."""


# This Callable annotation would cause a type error before Python 3.9.2, see
# https://docs.python.org/3/whatsnew/3.9.html#notable-changes-in-python-3-9-2.
if sys.version_info >= (3, 9, 2):
  ConvertFlagValuesFunction: TypeAlias = Callable[
      [PrimitiveFlagValue], PrimitiveFlagValue
  ]
  ConvertJsonFunction: TypeAlias = Callable[
      [NestedStrDict, Optional[str]], NestedStrDict
  ]
  ConvertStatusFunction: TypeAlias = Callable[[str, str, str], str]
  MatchOutputFunction: TypeAlias = Callable[[str], bool]
else:
  ConvertFlagValuesFunction: TypeAlias = Callable
  ConvertJsonFunction: TypeAlias = Callable
  ConvertStatusFunction: TypeAlias = Callable
  MatchOutputFunction: TypeAlias = Callable


class FlagMapping:
  """Defines how to create a gcloud command flag from a bq flag.

  For example this would return True:

  FlagMapping(
      bq_name='httplib2_debuglevel',
      gcloud_name='log-http',
      bq_to_gcloud_mapper=lambda x: x > 0,
  ).bq_to_gcloud_mapper(1)
  """

  def __init__(
      self,
      bq_name: str,  # The name of the original bq flag.
      gcloud_name: str,  # The gcloud flag that is mapped to.
      bq_to_gcloud_mapper: Optional[ConvertFlagValuesFunction] = None,
  ):
    self.bq_name = bq_name
    self.gcloud_name = gcloud_name
    if bq_to_gcloud_mapper:
      self.bq_to_gcloud_mapper = bq_to_gcloud_mapper
    else:
      self.bq_to_gcloud_mapper = self.default_map_bq_value_to_gcloud_value

  # TODO(b/398206856): Simplify this function and validate behaviour.
  def default_map_bq_value_to_gcloud_value(
      self, bq_flag_value: PrimitiveFlagValue
  ) -> PrimitiveFlagValue:
    """Takes a bq flag value and returns the equivalent gcloud flag value."""
    if isinstance(bq_flag_value, bool):
      return bq_flag_value or False
    elif isinstance(bq_flag_value, int):
      return bq_flag_value
    else:
      return str(bq_flag_value)


class UnsupportedFlagMapping(FlagMapping):
  """Defines a bq global flag that is not supported in gcloud."""

  def __init__(
      self,
      bq_name: str,
      error_message: str,
  ):
    def raise_unsupported_flag_error(x: Union[str, bool]) -> Union[str, bool]:
      raise BigqueryGcloudDelegationUserError(error_message)

    super().__init__(bq_name, 'unsupported_flag', raise_unsupported_flag_error)  # pyrefly: ignore[bad-argument-type]


def _convert_to_gcloud_flags(
    flag_mappings: Dict[str, FlagMapping],
    bq_flags: Dict[str, PrimitiveFlagValue],
) -> Dict[str, PrimitiveFlagValue]:
  """Returns the equivalent gcloud flags for a set of bq flags.

  Args:
    flag_mappings: The flag mappings to use. For example, {'project_id':
      FlagMapping('project_id', 'project')}
    bq_flags: The bq flags that will be mapped. For example, {'project_id':
      'my_project'}

  Returns:
    The equivalent gcloud flags. For example,
    {'project': 'my_project'}
  """
  gcloud_flags = {}
  for bq_flag, bq_flag_value in bq_flags.items():
    if bq_flag not in flag_mappings:
      raise ValueError(f'Unsupported bq flag: {bq_flag}')
    flag_mapper = flag_mappings[bq_flag]
    gcloud_flags[flag_mapper.gcloud_name] = flag_mapper.bq_to_gcloud_mapper(
        bq_flag_value
    )
  return gcloud_flags


class CommandMapping:
  """Stores the configuration to map a BQ CLI command to gcloud.

  This class does not include the global flags. These are handled at a higher
  level in the system.

  Example usage:

  CommandMapping(
      resource='datasets',
      bq_command='ls',
      gcloud_command=['alpha', 'bq', 'datasets', 'list'],
      flag_mapping_list=[
          FlagMapping(
              bq_name='max_results',
              gcloud_name='limit',
          ),
      ],
  ).get_gcloud_command_minus_global_flags(
      bq_format='pretty',
      bq_command_flags={'max_results': 5},
  )

  Results in:
  ['alpha', 'bq', 'datasets', 'list', '--format=table[box]', '--limit=5']
  """

  def __init__(
      self,
      resource: str,
      bq_command: str,
      gcloud_command: List[str],
      flag_mapping_list: Optional[List[FlagMapping]] = None,
      table_projection: Optional[str] = None,
      csv_projection: Optional[str] = None,
      json_mapping: Optional[ConvertJsonFunction] = None,
      status_mapping: Optional[ConvertStatusFunction] = None,
      synchronous_progress_message_matcher: Optional[
          MatchOutputFunction
      ] = None,
      print_resource: bool = True,
      no_prompts: bool = False,
  ):
    """Initializes the CommandMapping.

    Args:
      resource: The resource this command targets. For example, 'datasets'.
      bq_command: The bq command to map. For example, 'ls'.
      gcloud_command: The gcloud command that will be mapped to. For example,
        ['alpha', 'bq', 'datasets', 'list'].
      flag_mapping_list: The flag mappings for this command. For example,
        [FlagMapping('max_results', 'limit')]
      table_projection: An optional projection to use for the command when a
        table is displayed. For example:
        'datasetReference.datasetId:label=datasetId'.
      csv_projection: An optional projection to use for the command when the
        output is in csv format. For example:
        'datasetReference.datasetId:label=datasetId'.
      json_mapping: A function to map the json output from gcloud to bq. For
        example, lambda x: {'kind': 'bigquery#project', 'id': x['projectId']}
      status_mapping: A function to map the status output from gcloud to bq. For
        example, lambda orig, id, project: f'Dataset {project}:{id} deleted.'
      synchronous_progress_message_matcher: A function to match a progress
        message from gcloud when running a synchronous command. For example,
        lambda message: 'Waiting for job' in message.
      print_resource: If the command also prints the resource it is operating
        on. For example, 'ls' will list resources but 'rm' usually prints status
        and not the resource.
      no_prompts: Some commands need a prompt to be disabled when they're run
        and usually, the BQ CLI code flow will have done this already. For
        example, the when `bq rm -d` is run, the BQ CLI will prompt the user
        before deleting the dataset, so the gcloud prompt is not needed.
    """
    self.resource = resource
    self.bq_command = bq_command
    self.gcloud_command = gcloud_command
    self.flag_mapping_list = flag_mapping_list or []
    self._flag_mappings: Dict[str, FlagMapping] = None  # pyrefly: ignore[bad-assignment]
    self.table_projection = table_projection
    self.csv_projection = csv_projection
    self.json_mapping = json_mapping if json_mapping else lambda x, _: x
    if status_mapping:
      self.status_mapping = status_mapping
    else:
      self.status_mapping = lambda original_status, _, __: original_status
    if synchronous_progress_message_matcher:
      self.synchronous_progress_message_matcher = (
          synchronous_progress_message_matcher
      )
    else:
      self.synchronous_progress_message_matcher = lambda _: False
    self.print_resource = print_resource
    self.no_prompts = no_prompts

  @property
  def flag_mappings(self) -> Dict[str, FlagMapping]:
    """Returns the command flag mappings as a dictionary."""
    if not self._flag_mappings:
      self._flag_mappings = {}
      for flag_mapping in self.flag_mapping_list:
        self._flag_mappings[flag_mapping.bq_name] = flag_mapping
    return self._flag_mappings

  def _add_fields_to_format(
      self,
      prefix: str,
      labels: Optional[str] = None,
  ) -> str:
    """Returns the format from the map."""
    if labels:
      return f'{prefix}({labels})'
    else:
      return prefix

  def get_gcloud_format(self, bq_format: Optional[str]) -> str:
    """Returns the gcloud format for the given bq format."""
    # TODO(b/355324165): Update the note on what happens when there is no flag,
    # after we have better testing from the gcloud delegator.
    if not bq_format or bq_format == 'pretty' or bq_format == 'sparse':
      return self._add_fields_to_format('table[box]', self.table_projection)
    elif 'json' in bq_format:
      return 'json'
    elif 'csv' in bq_format:
      return self._add_fields_to_format('csv', self.csv_projection)
    else:
      raise ValueError(f'Unsupported format: {bq_format}')

  def _get_gcloud_flags(
      self,
      bq_flags: Dict[str, PrimitiveFlagValue],
  ) -> Dict[str, PrimitiveFlagValue]:
    """Returns the gcloud flags for the given bq flags."""
    return _convert_to_gcloud_flags(self.flag_mappings, bq_flags)

  def get_gcloud_command_minus_global_flags(
      self,
      bq_format: Optional[str],
      bq_command_flags: Dict[str, str],
      identifier: Optional[str] = None,
  ) -> List[str]:
    """Returns the gcloud command to use for the given bq command.

    Args:
      bq_format: The `format` flag from the BQ CLI (eg. 'json').
      bq_command_flags: The flags for this BQ command that will be mapped. For
        example, {'max_results': 5}
      identifier: An optional identifier of the resource this command will
        operate on.

    Returns:
      The equivalent gcloud command array with the leading 'gcloud' removed,
      with the format flag and command flags but no global flags. For example,
      ['alpha', 'bq', 'datasets', 'list', '--format=json', '--limit=5']
    """
    gcloud_command: List[str] = self.gcloud_command.copy()
    # If the resource is not being printed then don't add the format flag.
    if self.print_resource:
      gcloud_format = self.get_gcloud_format(bq_format)
      gcloud_command.append(
          f'--format={gcloud_format}',
      )
    gcloud_command.extend(
        _flatten_flag_dictionary(self._get_gcloud_flags(bq_command_flags))  # pyrefly: ignore[bad-argument-type]
    )
    if self.no_prompts:
      gcloud_command.append('--quiet')
    if identifier:
      gcloud_command.append(identifier)
    return gcloud_command


class GcloudCommandGenerator:
  """Generates a gcloud command from a bq command."""

  def __init__(
      self,
      command_mappings: List[CommandMapping],
      global_flag_mappings: List[FlagMapping],
  ):
    self._command_mapping_list = command_mappings
    self._global_flag_mapping_list = global_flag_mappings
    self._command_dict: Optional[Dict[str, Dict[str, CommandMapping]]] = None
    self._global_flag_dict: Optional[Dict[str, FlagMapping]] = None

  @property
  def command_dict(self) -> Dict[str, Dict[str, CommandMapping]]:
    """Returns the commands as a map of resource to bq command to delegator."""
    if not self._command_dict:
      self._command_dict = {}
      for command_mapping in self._command_mapping_list:
        if command_mapping.resource not in self._command_dict:
          self._command_dict[command_mapping.resource] = {}
        resource_to_commands = self._command_dict[command_mapping.resource]
        if command_mapping.bq_command in resource_to_commands:
          raise ValueError(
              f'Duplicate bq command: {command_mapping.bq_command}'
          )
        resource_to_commands[command_mapping.bq_command] = command_mapping
    return self._command_dict

  @property
  def global_flag_dict(self) -> Dict[str, FlagMapping]:
    if not self._global_flag_dict:
      self._global_flag_dict = {}
      for flag_mapping in self._global_flag_mapping_list:
        bq_flag = flag_mapping.bq_name
        if bq_flag in self._global_flag_dict:
          raise ValueError(f'Duplicate bq flag: {bq_flag}')
        self._global_flag_dict[bq_flag] = flag_mapping
    return self._global_flag_dict

  def map_to_gcloud_global_flags(
      self, bq_global_flags: Dict[str, PrimitiveFlagValue]
  ) -> Dict[str, PrimitiveFlagValue]:
    """Returns the equivalent gcloud global flags for a set of bq flags.

    In the Args and Returns below, this `GcloudCommandGenerator` is used:

    GcloudCommandGenerator(
      command_mappings=[],
      global_flag_mappings=[
        FlagMapping(
            bq_name='project_id',
            gcloud_name='project'),
        FlagMapping(
            bq_name='httplib2_debuglevel',
            gcloud_name='log-http', lambda x: x > 0)
    ])

    Args:
      bq_global_flags: The bq flags that will be mapped. For example,
        {'project_id': 'my_project', 'httplib2_debuglevel': 1}

    Returns:
      The equivalent gcloud flags. For example,
      {'project': 'my_project', 'log-http': True}
    """
    return _convert_to_gcloud_flags(self.global_flag_dict, bq_global_flags)

  def get_command_mapping(
      self, resource: str, bq_command: str
  ) -> CommandMapping:
    """Returns the gcloud delegator for the given resource and bq command."""
    # Fail fast if there is no CommandMapping.
    return self.command_dict[resource][bq_command]

  def get_gcloud_command(
      self,
      resource: str,
      bq_command: str,
      bq_global_flags: Dict[str, str],
      bq_command_flags: Dict[str, str],
      identifier: Optional[str] = None,
  ) -> List[str]:
    """Returns the gcloud command to use for the given bq command.

    As an example usage:

    GcloudCommandGenerator(
      command_mappings=[CommandMapping(
        resource='datasets',
        bq_command='ls',
        gcloud_command=['alpha', 'bq', 'datasets', 'list'],
        flag_mapping_list=[
            FlagMapping(
                bq_name='max_results',
                gcloud_name='limit',
            ),
      ],
      flag_mappings=[
        FlagMapping(
            bq_name='project_id',
            gcloud_name='project'),
    ]).get_gcloud_command(
        resource='datasets',
        bq_command='ls',
        bq_global_flags={'project_id': 'bigquery-cli-e2e', 'format': 'pretty'},
        bq_command_flags={'max_results': 5},
    )

    Will return:

    ['--project=bigquery-cli-e2e', 'alpha', 'bq', 'datasets', 'list',
    '--format=json', '--limit=5']

    Args:
      resource: The resource the command is being run on, named to align with
        `gcloud` commands. For example, 'jobs' or 'datasets'.
      bq_command: The bq command to run. For example, 'ls' or 'show'.
      bq_global_flags: The BQ CLI global flags for the command.
      bq_command_flags: The BQ CLI command flags for the command.
      identifier: The identifier of the resource to act on.

    Returns:
      The gcloud command to run as an array of strings, minus the leading
      'gcloud'. This can be parsed directly into
      `gcloud_runner.run_gcloud_command`.
    """
    delegator = self.get_command_mapping(resource, bq_command)
    if not delegator:
      raise ValueError(f'Unsupported bq command: {bq_command}')

    # TODO(b/355324165): Revisit how the format flag is passed.
    # The format flag is handled separately so filter it out.
    filtered_global_flags = bq_global_flags.copy()
    bq_format = filtered_global_flags.pop('format', 'sparse')

    gcloud_global_flags: List[str] = _flatten_flag_dictionary(
        self.map_to_gcloud_global_flags(filtered_global_flags)  # pyrefly: ignore[bad-argument-type]
    )

    return (
        gcloud_global_flags
        + delegator.get_gcloud_command_minus_global_flags(
            bq_format=bq_format,
            bq_command_flags=bq_command_flags,
            identifier=identifier,
        )
    )
