# -*- coding: utf-8 -*- #
# Copyright 2026 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.
"""Implements command to translate a SQL query."""

import functools
import json
import os
import typing

from apitools.base.py import encoding
from googlecloudsdk.api_lib.bq import util as api_util
from googlecloudsdk.api_lib.util import apis
from googlecloudsdk.api_lib.util import waiter
from googlecloudsdk.calliope import arg_parsers
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.bq import command_utils
from googlecloudsdk.core import exceptions
from googlecloudsdk.core import log
from googlecloudsdk.core import properties
from googlecloudsdk.core import resources
from googlecloudsdk.core import yaml
from googlecloudsdk.core.console import console_io
from googlecloudsdk.core.util import files


@functools.cache
def _get_dialect_registry():
  """Gets dialect registry: input_dialects, output_dialects, dialect_pairs.

  The mapping is loaded from a JSON resource file.

  Returns:
    A dict with dialect registry.

  Raises:
    exceptions.Error: If the dialect registry fails to load.
  """
  try:
    file_path = os.path.join(os.path.dirname(__file__), 'dialect_registry.json')
    registry = json.loads(files.ReadFileContents(file_path))
    return registry
  except Exception as e:
    raise exceptions.Error(f'Failed to load dialect registry: {e}')


def _get_task_type(source_dialect: str, target_dialect: str) -> str:
  """Returns the migration task type based on the source dialect."""
  registry = _get_dialect_registry()
  source_dialect_map = {}
  for dialect in registry.get('input_dialects', []):
    legacy_name = dialect.get('legacy_batch_name')
    name = dialect.get('name')
    if legacy_name:
      if name:
        source_dialect_map[name.lower()] = legacy_name
      source_dialect_map[legacy_name.lower()] = legacy_name
    elif name:
      source_dialect_map[name.lower()] = name
  source_legacy_batch_name = source_dialect_map.get(
      source_dialect.lower(), source_dialect
  )
  target_dialect_map = {}
  for dialect in registry.get('output_dialects', []):
    legacy_name = dialect.get('legacy_batch_name')
    name = dialect.get('name')
    if legacy_name:
      if name:
        target_dialect_map[name.lower()] = legacy_name
      target_dialect_map[legacy_name.lower()] = legacy_name
    elif name:
      target_dialect_map[name.lower()] = name
  target_legacy_batch_name = target_dialect_map.get(
      target_dialect.lower(), target_dialect
  )
  task_type = (
      f'{source_legacy_batch_name}2{target_legacy_batch_name}_Translation'
  )
  valid_task_type = False
  for pair in registry['dialect_pairs']:
    if task_type in pair.get('legacy_batch_name', []):
      valid_task_type = True
      break
  if not valid_task_type:
    raise exceptions.Error(
        f'Translation from {source_dialect} to {target_dialect} is not'
        ' supported.'
    )
  return task_type


def _build_translation_details(
    messages,
    query,
    explanation_output_file=None,
    translation_config_files=None,
    metadata_gcs_uri=None,
    source_ddl_output_file=None,
):
  """Builds the translation details message for the migration task."""
  target_return_literals = ['sql/query.sql']
  target_types = ['sql']

  source_target_mapping = [
      messages.GoogleCloudBigqueryMigrationV2SourceTargetMapping(
          sourceSpec=messages.GoogleCloudBigqueryMigrationV2SourceSpec(
              literal=messages.GoogleCloudBigqueryMigrationV2Literal(
                  literalString=query, relativePath='query.sql'
              )
          )
      )
  ]

  if source_ddl_output_file:
    target_return_literals.append('source_sql_metadata_generation_suggestion/')
    if 'suggestion' not in target_types:
      target_types.append('suggestion')

  if explanation_output_file:
    target_return_literals.append(
        'translation_explanation_suggestion/query.sql'
    )
    if 'suggestion' not in target_types:
      target_types.append('suggestion')
    ai_config = (
        'translation_rules:\n  - suggestion_type: "TRANSLATION_EXPLANATION"'
    )
    source_target_mapping.append(
        messages.GoogleCloudBigqueryMigrationV2SourceTargetMapping(
            sourceSpec=messages.GoogleCloudBigqueryMigrationV2SourceSpec(
                literal=messages.GoogleCloudBigqueryMigrationV2Literal(
                    literalString=ai_config,
                    relativePath='gcloud_default_inline.ai_config.yaml',
                )
            )
        )
    )

  if translation_config_files:
    for config_file in translation_config_files:
      filename = os.path.basename(config_file)
      if not filename.endswith('.config.yaml'):
        raise exceptions.Error(
            f'Translation config file "{config_file}" must end with'
            ' ".config.yaml".'
        )
      content = files.ReadFileContents(config_file)
      source_target_mapping.append(
          messages.GoogleCloudBigqueryMigrationV2SourceTargetMapping(
              sourceSpec=messages.GoogleCloudBigqueryMigrationV2SourceSpec(
                  literal=messages.GoogleCloudBigqueryMigrationV2Literal(
                      literalString=content,
                      relativePath=filename,
                  )
              )
          )
      )

  if metadata_gcs_uri:
    source_target_mapping.append(
        messages.GoogleCloudBigqueryMigrationV2SourceTargetMapping(
            sourceSpec=messages.GoogleCloudBigqueryMigrationV2SourceSpec(
                gcsFilePath=metadata_gcs_uri
            )
        )
    )

  return messages.GoogleCloudBigqueryMigrationV2TranslationDetails(
      targetReturnLiterals=target_return_literals,
      targetTypes=target_types,
      sourceTargetMapping=source_target_mapping,
  )


def _build_migration_workflow(
    messages,
    query,
    source_dialect,
    target_dialect,
    explanation_output_file=None,
    translation_config_files=None,
    metadata_gcs_uri=None,
    source_ddl_output_file=None,
):
  """Builds the migration workflow message and returns it with the task type."""
  task_type = _get_task_type(source_dialect, target_dialect)
  translation_details = _build_translation_details(
      messages,
      query,
      explanation_output_file,
      translation_config_files,
      metadata_gcs_uri,
      source_ddl_output_file,
  )

  task = messages.GoogleCloudBigqueryMigrationV2MigrationTask(
      type=task_type, translationDetails=translation_details
  )

  workflow_tasks_value = messages.GoogleCloudBigqueryMigrationV2MigrationWorkflow.TasksValue(
      additionalProperties=[
          messages.GoogleCloudBigqueryMigrationV2MigrationWorkflow.TasksValue.AdditionalProperty(
              key='translation_task', value=task
          )
      ]
  )
  workflow = messages.GoogleCloudBigqueryMigrationV2MigrationWorkflow(
      tasks=workflow_tasks_value
  )
  return workflow, task_type


def _parse_task_result(wait_response):
  """Parses the workflow wait response into component task details."""
  translated_sql = None
  explanation = None
  generated_ddl = None
  result_task = None
  console_uri = ''
  report_log_messages = []

  if wait_response.tasks and wait_response.tasks.additionalProperties:
    translation_task_prop = next(
        (
            p
            for p in wait_response.tasks.additionalProperties
            if p.key == 'translation_task'
        ),
        None,
    )
    if translation_task_prop:
      result_task = translation_task_prop.value
    if (
        result_task
        and result_task.taskResult
        and result_task.taskResult.translationTaskResult
    ):
      literals = result_task.taskResult.translationTaskResult.translatedLiterals
      for literal in literals:
        if literal.relativePath == 'sql/query.sql':
          translated_sql = literal.literalString
        elif (
            literal.relativePath
            == 'translation_explanation_suggestion/query.sql'
        ):
          explanation = literal.literalString
        elif (
            literal.relativePath == 'source_sql_metadata_generation_suggestion/'
        ):
          generated_ddl = literal.literalString
      console_uri = result_task.taskResult.translationTaskResult.consoleUri
      report_log_messages = (
          result_task.taskResult.translationTaskResult.reportLogMessages
      )

  return (
      result_task,
      translated_sql,
      explanation,
      generated_ddl,
      console_uri,
      report_log_messages,
  )


def _format_log_result(
    wait_response, task_type, result_task, console_uri, report_log_messages
):
  """Formats the log result dictionary from the task components."""
  result = {
      'name': wait_response.name,
      'type': task_type,
      'state': wait_response.state.name if wait_response.state else '',
      'consoleUri': console_uri,
      'reportLogMessages': (
          [encoding.MessageToDict(msg) for msg in report_log_messages]
          if report_log_messages
          else []
      ),
  }

  if result_task:
    if result_task.processingError:
      result['processingError'] = encoding.MessageToDict(
          result_task.processingError
      )
    if result_task.resourceErrorDetails:
      result['resourceErrorDetails'] = [
          encoding.MessageToDict(msg)
          for msg in result_task.resourceErrorDetails
      ]

  return result


def _handle_task_failure(wait_response, result_task):
  """Checks the task state for failure and raises an exception if needed."""
  if result_task and result_task.state and result_task.state.name == 'FAILED':
    log.error('Workflow wait_response: %s', wait_response)
    if result_task.processingError:
      log.error('Task processing error: %s', result_task.processingError)
      error_details = []
      if (
          hasattr(result_task.processingError, 'message')
          and result_task.processingError.message
      ):
        error_details.append(result_task.processingError.message)
      if (
          hasattr(result_task.processingError, 'reason')
          and result_task.processingError.reason
      ):
        error_details.append(result_task.processingError.reason)
      error_str = ' '.join(error_details) if error_details else 'Unknown error'
      log.err.Print(f'Task processing error: {error_str}')
    if result_task.resourceErrorDetails:
      log.error('Task resource errors: %s', result_task.resourceErrorDetails)
    raise exceptions.Error('Failed to retrieve the translated SQL query.')


def _process_workflow_result(wait_response, task_type, args):
  """Processes the completed workflow and handles output/logging."""
  (
      result_task,
      translated_sql,
      explanation,
      generated_ddl,
      console_uri,
      report_log_messages,
  ) = _parse_task_result(wait_response)

  result = _format_log_result(
      wait_response, task_type, result_task, console_uri, report_log_messages
  )

  if args.translation_log_file:
    files.WriteFileContents(args.translation_log_file, yaml.dump(result))
    log.status.Print(
        f'Successfully saved translation logs to {args.translation_log_file}.'
    )

  _handle_task_failure(wait_response, result_task)

  if not translated_sql:
    raise exceptions.Error('Failed to retrieve the translated SQL query.')

  if args.output_file and args.output_file != '-':
    files.WriteFileContents(args.output_file, translated_sql)
    log.status.Print(
        f'Successfully translated and saved to {args.output_file}.'
    )
  else:
    log.out.Print('Translated query:')
    log.out.Print(translated_sql)

  if args.explanation_output_file and explanation:
    files.WriteFileContents(args.explanation_output_file, explanation)
    log.status.Print(
        'Successfully saved translation explanation to'
        f' {args.explanation_output_file}.'
    )

  if args.source_ddl_output_file and generated_ddl:
    files.WriteFileContents(args.source_ddl_output_file, generated_ddl)
    log.status.Print(
        f'Successfully saved source DDL to {args.source_ddl_output_file}.'
    )

  if args.translation_log_file:
    # Return None to avoid duplicating output in standard terminal output.
    return None

  log.out.Print('\nTranslation logs:')
  return result


@base.UniverseCompatible
@base.ReleaseTracks(base.ReleaseTrack.ALPHA)
class Translate(base.Command):
  """Translate a SQL query from a source dialect to BigQuery."""

  detailed_help = {
      'brief': 'Translate a SQL query from a source dialect to BigQuery.',
      'DESCRIPTION': """\
Translate a SQL query from a source dialect to BigQuery.

The command reads a SQL query from standard input (or an input file) and
prints the translated query to standard output (or an output file).
""",
      'EXAMPLES': """\
To translate a Snowflake query from stdin, run:

  $ echo 'SELECT * FROM test.my_table;' | {command} --source-dialect=SNOWFLAKE --location=us

To translate a Snowflake query from a file and save the output and logs to files, run:

  $ {command} \\
      --source-dialect=SNOWFLAKE \\
      --location=us \\
      --project=my-project \\
      --input-file=input.sql \\
      --output-file=output.sql \\
      --translation-log-file=translation_logs.yaml
""",
  }

  @staticmethod
  def Args(parser):
    parser.add_argument(
        '--source-dialect',
        help=(
            'Source dialect of the query. See supported dialects in'
            ' https://docs.cloud.google.com/bigquery/docs/batch-sql-translator#supported_sql_dialects'
        ),
        required=True,
    )
    parser.add_argument(
        '--target-dialect',
        help=(
            'Target dialect of the query. See supported dialects in'
            ' https://docs.cloud.google.com/bigquery/docs/batch-sql-translator#supported_sql_dialects'
        ),
        required=True,
    )
    parser.add_argument(
        '--location',
        help='Google Cloud Storage location to use for the translation.',
        required=True,
    )
    parser.add_argument(
        '--output-file',
        help=(
            'File to which to write the translated SQL. Writes to stdout if'
            " not provided or if set to ``-''."
        ),
        required=False,
    )
    parser.add_argument(
        '--input-file',
        help=(
            'File from which to read the SQL query to translate. Can be piped'
            ' through stdin as well.'
        ),
        required=False,
    )
    parser.add_argument(
        '--translation-log-file',
        help='File to which to write the translation logs.',
        required=False,
    )
    parser.add_argument(
        '--explanation-output-file',
        help='File to which to write the translation explanation.',
        required=False,
    )
    parser.add_argument(
        '--metadata-gcs-uri',
        help='Cloud Storage URI for the metadata zip file.',
        required=False,
    )
    parser.add_argument(
        '--translation-config-files',
        type=arg_parsers.ArgList(),
        metavar='CONFIG_FILE',
        help=(
            'A comma-separated list of paths to YAML configuration files. File'
            ' names must end with .config.yaml. See'
            ' https://docs.cloud.google.com/bigquery/docs/config-yaml-translation'
            ' for more details.'
        ),
        required=False,
    )
    parser.add_argument(
        '--source-ddl-output-file',
        help='File to which to write the generated source DDL.',
        required=False,
    )

  def Run(self, args):
    query = console_io.ReadFromFileOrStdin(args.input_file or '-', binary=False)

    if not query:
      raise exceptions.Error('No input query provided.')

    client = api_util.GetMigrationApiClient()
    client.additional_http_headers[api_util.TOOL_TAG_HEADER] = (
        api_util.GCLOUD_TOOL_TAG
    )
    messages = typing.cast(
        typing.Any, apis.GetMessagesModule('bigquerymigration', 'v2')
    )
    migration_service = client.projects_locations_workflows

    project = args.project or properties.VALUES.core.project.Get(required=True)
    location = args.location

    workflow, task_type = _build_migration_workflow(
        messages,
        query,
        args.source_dialect,
        args.target_dialect,
        explanation_output_file=args.explanation_output_file,
        translation_config_files=args.translation_config_files,
        metadata_gcs_uri=args.metadata_gcs_uri,
        source_ddl_output_file=args.source_ddl_output_file,
    )

    request_type = api_util.GetMigrationApiMessage(
        'BigquerymigrationProjectsLocationsWorkflowsCreateRequest'
    )
    request = request_type(
        parent=f'projects/{project}/locations/{location}',
        googleCloudBigqueryMigrationV2MigrationWorkflow=workflow,
    )

    response = migration_service.Create(request)

    workflow_ref = resources.REGISTRY.ParseRelativeName(
        response.name,
        collection='bigquerymigration.projects.locations.workflows',
    )
    poller = command_utils.BqMigrationWorkflowPoller(migration_service)

    wait_response = waiter.WaitFor(
        poller=poller,
        operation_ref=workflow_ref,
        message='Running translation workflow [{}]'.format(response.name),
    )

    return _process_workflow_result(wait_response, task_type, args)
