# -*- coding: utf-8 -*- #
# Copyright 2026 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.
"""`gcloud dataplex dbt metadata-jobs create` command."""

from __future__ import annotations

import os
import uuid

from apitools.base.py import encoding
from apitools.base.py import exceptions as apitools_exceptions
from googlecloudsdk.api_lib.cloudresourcemanager import projects_api
from googlecloudsdk.api_lib.dataplex import dbt_metadata_job as dbt_job_lib
from googlecloudsdk.api_lib.dataplex import entry_group as entry_group_lib
from googlecloudsdk.api_lib.dataplex import metadata_job as metadata_job_lib
from googlecloudsdk.api_lib.dataplex import util as dataplex_util
from googlecloudsdk.api_lib.storage import storage_api
from googlecloudsdk.api_lib.storage import storage_util
from googlecloudsdk.api_lib.util import exceptions as gcloud_exception
from googlecloudsdk.api_lib.util import waiter
from googlecloudsdk.calliope import base
from googlecloudsdk.calliope import parser_arguments
from googlecloudsdk.calliope import parser_extensions
from googlecloudsdk.command_lib.dataplex import resource_args
from googlecloudsdk.command_lib.dataplex.dbt import bigquery_location as bq_loc
from googlecloudsdk.command_lib.dataplex.dbt import transform as dbt_transform
from googlecloudsdk.command_lib.projects import util as projects_util
from googlecloudsdk.command_lib.util.apis import arg_utils
from googlecloudsdk.core import exceptions
from googlecloudsdk.core import log
from googlecloudsdk.core import resources
from googlecloudsdk.core.util import files

_JSONL_FILENAME = 'dbt_metadata.jsonl'

# Metadata job states that mean the import is no longer running.
_TERMINAL_STATES = frozenset(
    ['SUCCEEDED', 'SUCCEEDED_WITH_ERRORS', 'FAILED', 'CANCELED']
)


class _ImportJobPoller(waiter.OperationPoller):
  """Polls a metadata import job until it reaches a terminal state.

  The metadataJobs.create operation completes when the job is accepted, not when
  the import finishes, so the job's own status has to be polled to learn the
  real outcome (and entry counts).
  """

  def __init__(self, jobs_service, messages):
    self._jobs_service = jobs_service
    self._messages = messages

  def IsDone(self, job):
    state = job.status.state if job and job.status else None
    return str(state) in _TERMINAL_STATES if state else False

  def Poll(self, job_name):
    return self._jobs_service.Get(
        self._messages.DataplexProjectsLocationsMetadataJobsGetRequest(
            name=job_name
        )
    )

  def GetResult(self, job):
    return job


@base.Hidden
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.ALPHA)
class Create(base.Command):
  """Transform dbt-core artifacts and import them into Dataplex Catalog.

  This command reads the JSON artifacts produced by dbt-core (manifest.json,
  catalog.json, run_results.json, sources.json) from a local directory,
  transforms them into the Dataplex metadata import format, uploads the result
  to Cloud Storage, and triggers a Dataplex metadata import job that ingests
  the metadata into the Knowledge Catalog.

  Only the entry group that receives the dbt entries must exist in the caller's
  project beforehand. The caller must also be able to USE the dbt connector
  types (dataplex.aspectTypes.use / the dbt-connector-types alternate-use
  permission).

  The Metadata Job ID identifies the import run and, if provided, must:
   * Contain only lowercase letters, numbers, and hyphens.
   * Start with a letter and end with a number or a letter.
   * Be 1-63 characters and unique within the project / location.
  """

  detailed_help = {
      'EXAMPLES': (
          """\
          To transform the dbt artifacts in the current directory and import
          them into entry group `dbt-metadata-ingestion` in project
          `my-project`, location `us-central1`, run:

            $ {command} my-dbt-import --project=my-project \
                --location=us-central1 \
                --artifacts-path=. \
                --entry-group=dbt-metadata-ingestion \
                --storage-uri=gs://my-bucket/dbt-imports/

          To only build and upload the JSONL and validate the job without
          ingesting, add `--validate-only`.
          """
      ),
  }

  @staticmethod
  def Args(parser: parser_arguments.ArgumentInterceptor) -> None:
    resource_args.AddMetadataJobResourceArg(parser, 'to create.')
    parser.add_argument(
        '--artifacts-path',
        default='.',
        help="""Local path to the dbt-core artifacts. May point at the dbt
        project root (the `target/` subdirectory is detected automatically) or
        directly at the directory containing manifest.json. Defaults to the
        current working directory.""",
    )
    parser.add_argument(
        '--storage-uri',
        required=True,
        help="""Cloud Storage URI prefix (gs://bucket/path/) the transformed
        JSONL is uploaded to and the import job reads from. The caller must have
        write access and the Dataplex service agent must have read access.""",
    )
    parser.add_argument(
        '--entry-group',
        default='dbt-metadata-ingestion',
        help="""Short ID of the entry group that receives the dbt entries. Must
        already exist in the project / location.""",
    )
    parser.add_argument(
        '--connector-types-project',
        hidden=True,
        help="""Overrides the project that owns the 1P dbt aspect/entry types.
        Defaults automatically; for internal/testing use only.""",
    )
    parser.add_argument(
        '--system-types-project',
        hidden=True,
        help="""Overrides the project that owns the core 1P types the dbt types
        depend on (the `contacts` aspect type and the dbt entry link types).
        Defaults automatically; for internal/testing use only.""",
    )
    parser.add_argument(
        '--import-entry-sync-mode',
        choices={
            'FULL': (
                """All entries in the job scope are synced; entries absent
                    from the import file are deleted."""
            ),
            'INCREMENTAL': (
                """Only entries present in the import file are
                    modified."""
            ),
        },
        type=arg_utils.ChoiceToEnumName,
        default='FULL',
        help='Entry sync mode for the import job.',
    )
    parser.add_argument(
        '--import-aspect-sync-mode',
        choices={
            'FULL': """All aspects in the job scope are synced.""",
            'INCREMENTAL': (
                """Only aspects present in the import file are
                    modified."""
            ),
        },
        type=arg_utils.ChoiceToEnumName,
        default='INCREMENTAL',
        help='Aspect sync mode for the import job.',
    )
    parser.add_argument(
        '--include-entry-links',
        action='store_true',
        default=True,
        help="""Also emit EntryLink records capturing dbt lineage and semantic
        relationships (depends-on-lineage-imported, represents, depends-on-imported, etc.).""",
    )
    parser.add_argument(
        '--skip-bigquery-link',
        action='store_true',
        default=False,
        help="""Skip `represents` links (dbt node -> physical BigQuery
        table entry). Otherwise a `represents` link is emitted for each
        materialized dbt node (model/seed/snapshot) whose BigQuery dataset lives
        in the import location (`--location`); links can only reference
        @bigquery entries in that same region, so datasets in another region are
        skipped automatically. Use this flag when the BigQuery tables are not
        cataloged in Dataplex.""",
    )
    parser.add_argument(
        '--validate-only',
        action='store_true',
        default=False,
        help="""Build and upload the JSONL and validate the metadata job, but
        don't actually ingest.""",
    )
    base.ASYNC_FLAG.AddToParser(parser)

  @gcloud_exception.CatchHTTPErrorRaiseHTTPException(
      'Status code: {status_code}. {status_message}.'
  )
  def Run(self, args: parser_extensions.Namespace) -> None:
    metadata_job = args.CONCEPTS.metadata_job.Parse()
    parent = metadata_job.Parent().RelativeName()
    project_id = metadata_job.projectsId
    location = metadata_job.locationsId
    metadata_job_id = self._GetMetadataJobId(metadata_job)

    # Entry names / entry-group refs use the project NUMBER.
    project_number = self._GetProjectNumber(project_id)

    # The 1P dbt types live in env-specific system projects at `global`:
    # dbt aspect/entry types in the connector project, contacts + entry link
    # types in the core system project.
    connector_types_project = dbt_job_lib.ResolveConnectorTypesProject(
        args.connector_types_project
    )
    system_types_project = dbt_job_lib.ResolveSystemTypesProject(
        args.system_types_project
    )
    # The 1P dbt types always live at the `global` location.
    types_location = 'global'
    log.status.Print(
        'Using dbt types from [{0}] and core types from [{1}] (location {2}).'
        .format(connector_types_project, system_types_project, types_location)
    )

    # Fail fast if the target entry group is missing, before spending time
    # transforming artifacts and uploading them to GCS (otherwise this surfaces
    # only later, in the asynchronous import job).
    self._CheckEntryGroupExists(
        project_number, project_id, location, args.entry_group
    )

    # Resolve which datasets get represents (physical) links
    # (dbt node -> physical @bigquery table entry). Those links can only
    # reference @bigquery entries in the import location, so datasets in another
    # region are dropped.
    linkable_datasets = self._ResolveLinkableDatasets(args, location)

    # 1. Transform dbt artifacts into a JSONL import file in a temp dir.
    with files.TemporaryDirectory() as tmp_dir:
      local_jsonl = os.path.join(tmp_dir, _JSONL_FILENAME)
      summary = dbt_transform.GenerateImportFile(
          artifacts_path=args.artifacts_path,
          output_path=local_jsonl,
          eg_project=project_number,
          eg_project_id=project_id,
          eg_location=location,
          entry_group=args.entry_group,
          connector_types_project=connector_types_project,
          system_types_project=system_types_project,
          types_location=types_location,
          include_entry_links=args.include_entry_links,
          linkable_datasets=linkable_datasets,
      )
      log.status.Print(
          'Transformed dbt artifacts: {0} entries, {1} entry links.'.format(
              summary['entries'], summary['entry_links']
          )
      )

      # 2. Upload the JSONL under a per-job prefix (avoids stale-file
      #    duplicates) and point the import job at that prefix.
      storage_prefix = self._JobStoragePrefix(args.storage_uri, metadata_job_id)
      object_uri = storage_prefix + _JSONL_FILENAME
      log.status.Print('Uploading import file to {0} ...'.format(object_uri))
      storage_api.StorageClient().CopyFileToGCS(
          local_jsonl, storage_util.ObjectReference.FromUrl(object_uri)
      )

    # 3. Build and submit the import job referencing the dbt connector types.
    entry_link_types = None
    referenced_entry_scopes = None
    extra_aspect_types = None
    if args.include_entry_links:
      entry_link_types = dbt_transform.LinkTypeFqns(
          system_types_project, types_location
      )
      # Scope the caller's project plus any BigQuery projects that
      # represents (physical) links target, so those cross-entry references
      # resolve.
      referenced_entry_scopes = ['projects/{0}'.format(project_number)] + [
          'projects/{0}'.format(p) for p in summary.get('bigquery_projects', [])
      ]

    job = dbt_job_lib.GenerateImportMetadataJob(
        eg_project=project_number,
        eg_location=location,
        entry_group=args.entry_group,
        connector_types_project=connector_types_project,
        system_types_project=system_types_project,
        source_storage_uri=storage_prefix,
        entry_sync_mode=args.import_entry_sync_mode,
        aspect_sync_mode=args.import_aspect_sync_mode,
        entry_link_types=entry_link_types,
        referenced_entry_scopes=referenced_entry_scopes,
        extra_aspect_types=extra_aspect_types,
    )

    dataplex_client = dataplex_util.GetClientInstance()
    message = dataplex_util.GetMessageModule()
    create_req_op = dataplex_client.projects_locations_metadataJobs.Create(
        message.DataplexProjectsLocationsMetadataJobsCreateRequest(
            metadataJobId=metadata_job_id,
            parent=parent,
            googleCloudDataplexV1MetadataJob=job,
            validateOnly=args.validate_only,
        ),
    )

    if args.validate_only:
      log.status.Print('Validation complete.')
      return

    # Always surface the operation and the job ID.
    job_id = metadata_job_id or self._ServerGeneratedJobId(create_req_op)
    if job_id:
      log.status.Print(
          'Submitted dbt metadata import job [{0}] with operation [{1}].'
          .format(job_id, create_req_op.name)
      )
    else:
      log.status.Print(
          'Submitted dbt metadata import job with operation [{0}].'.format(
              create_req_op.name
          )
      )

    if getattr(args, 'async_', False):
      return

    # The create operation only confirms the job was accepted. Wait for it, then
    # poll the job itself until the import reaches a terminal state so we can
    # report the real outcome (and fail on a failed import).
    metadata_job_lib.WaitForOperation(create_req_op)
    if not job_id:
      # Without the id we can't address the job to poll; the create succeeded.
      log.status.Print(
          'dbt metadata import job created in [{0}].'.format(parent)
      )
      return
    result = self._WaitForImport(
        dataplex_client, message, '{0}/metadataJobs/{1}'.format(parent, job_id)
    )
    self._ReportImportOutcome(job_id, result)

  def _ResolveLinkableDatasets(
      self, args: parser_extensions.Namespace, location: str
  ) -> frozenset[tuple[str, str]] | None:
    """Returns the BigQuery datasets to emit represents links for.

    A represents link points a dbt node at its physical @bigquery table
    entry. The link is created in the caller's entry group at the import
    location; Dataplex only supports same-region entry links, and the @bigquery
    entry of a BigQuery table lives in the Dataplex region matching its dataset.
    So a link is only valid when the dataset is in the import location -- its
    region is not a choice, it is always the import location. A live BigQuery
    lookup is used only to drop datasets that demonstrably live elsewhere;
    callers who can't do that lookup or don't want these links pass
    --skip-bigquery-link.

    Args:
      args: the parsed command arguments.
      location: the import location (the entry group / metadata job region).

    Returns:
      The set of (project, dataset) pairs to emit links for, or None when
      represents links are disabled or nothing is co-located with the
      import location.
    """
    if not args.include_entry_links or args.skip_bigquery_link:
      return None
    datasets = dbt_transform.MaterializedBigQueryDatasets(args.artifacts_path)
    if not datasets:
      return None
    # Drop only datasets we can prove live in another region. Datasets we can't
    # read (no bigquery.datasets.get access, or not found) are kept
    # optimistically -- the import reports an unresolved @bigquery target as a
    # non-fatal per-link error.
    resolved = bq_loc.ResolveDatasetLocations(datasets)
    mismatched = {
        dataset: region
        for dataset, region in resolved.items()
        if region != location
    }
    linkable = frozenset(datasets - set(mismatched))
    if mismatched:
      log.warning(
          'Skipping represents links for {0} BigQuery dataset(s) not in '
          'the import location [{1}]: {2}. Entry links must be same-region, so '
          '@bigquery entries in another region cannot be linked; run the '
          'import in that region (--location) to link them.'.format(
              len(mismatched),
              location,
              ', '.join(
                  '{0}.{1} [{2}]'.format(project, dataset, region)
                  for (project, dataset), region in sorted(mismatched.items())
              ),
          )
      )
    return linkable or None

  def _CheckEntryGroupExists(
      self,
      project_number: str,
      project_id: str,
      location: str,
      entry_group: str,
  ) -> None:
    """Fails early with an actionable message if the entry group is absent.

    Only a genuine "not found" is treated as fatal here; any other error (e.g. a
    transient failure, or a permission check that Get is stricter about than the
    import job) is left to surface later rather than blocking the import on a
    best-effort pre-flight check.

    Args:
      project_number: project NUMBER owning the entry group (used in the name).
      project_id: project ID, for the actionable error message.
      location: Dataplex region of the entry group.
      entry_group: short id of the entry group.

    Raises:
      exceptions.Error: if the entry group does not exist.
    """
    name = 'projects/{0}/locations/{1}/entryGroups/{2}'.format(
        project_number, location, entry_group
    )
    try:
      entry_group_lib.GetEntryGroup(name)
    except apitools_exceptions.HttpNotFoundError as exc:
      raise exceptions.Error(
          'Entry group [{entry_group}] does not exist in project '
          '[{project_id}], location [{location}]. Create it first, e.g.:\n'
          '  gcloud dataplex entry-groups create {entry_group} '
          '--project={project_id} --location={location}\n'
          'then re-run this command.'.format(
              entry_group=entry_group,
              project_id=project_id,
              location=location,
          )
      ) from exc
    except apitools_exceptions.HttpError as exc:
      log.debug(
          'Ignoring non-404 error from entry group pre-flight check: %s', exc
      )

  def _GetMetadataJobId(self, metadata_job: resources.Resource) -> str | None:
    metadata_job_id = metadata_job.RelativeName().split('/')[-1]
    if metadata_job_id == resource_args.GENERATE_ID:
      return None
    return metadata_job_id

  def _ServerGeneratedJobId(self, operation) -> str | None:
    """Returns the metadata job id the server assigned to a create operation.

    The id is only unknown locally when it wasn't passed on the command line.
    The server records it as the operation's target resource path when the
    operation is created -- available even with --async -- so it can be surfaced
    without waiting for the import to finish. Returns None if it isn't there.

    Args:
      operation: The create long-running operation.
    """
    if operation is None or operation.metadata is None:
      return None
    target = encoding.MessageToPyValue(operation.metadata).get('target')
    return target.split('/')[-1] if target else None

  def _WaitForImport(self, dataplex_client, messages, job_name):
    """Polls the metadata job until the import reaches a terminal state.

    Args:
      dataplex_client: The Dataplex API client.
      messages: The Dataplex message module.
      job_name: The full metadata job resource name to poll.

    Returns:
      The finished GoogleCloudDataplexV1MetadataJob resource.
    """
    poller = _ImportJobPoller(
        dataplex_client.projects_locations_metadataJobs, messages
    )
    return waiter.WaitFor(
        poller,
        job_name,
        'Waiting for dbt metadata import job [{0}] to finish'.format(
            job_name.split('/')[-1]
        ),
    )

  def _ReportImportOutcome(self, job_id: str | None, result) -> None:
    """Reports the finished import job's state and entry counts.

    The create operation completing doesn't mean the import succeeded: the job
    can finish SUCCEEDED, SUCCEEDED_WITH_ERRORS, FAILED or CANCELED. Surface the
    real state (and the entry counts) rather than a blanket "created", and fail
    the command on a failed or canceled import.

    Args:
      job_id: The metadata job id, if known.
      result: The finished GoogleCloudDataplexV1MetadataJob resource.

    Raises:
      exceptions.Error: If the import job failed or was canceled.
    """
    label = 'dbt metadata import job'
    if job_id:
      label = '{0} [{1}]'.format(label, job_id)
    status = getattr(result, 'status', None)
    state = (
        str(status.state) if status and status.state else 'STATE_UNSPECIFIED'
    )
    message = status.message if status and status.message else ''
    suffix = ': {0}'.format(message) if message else ''

    if state in ('FAILED', 'CANCELED'):
      raise exceptions.Error('{0} {1}{2}'.format(label, state.lower(), suffix))
    if state == 'SUCCEEDED_WITH_ERRORS':
      log.warning('{0} completed with errors{1}'.format(label, suffix))
    elif state == 'SUCCEEDED':
      log.status.Print('{0} succeeded.'.format(label))
    else:
      log.status.Print('{0} finished.'.format(label))

    counts = self._ImportCounts(getattr(result, 'importResult', None))
    if counts:
      log.status.Print('  {0}'.format(counts))

  def _ImportCounts(self, import_result) -> str:
    """Returns a compact summary of import entry / link counts, or ''."""
    if import_result is None:
      return ''
    entries = ', '.join(
        '{0} {1}'.format(n, name)
        for name, n in (
            ('created', import_result.createdEntries),
            ('updated', import_result.updatedEntries),
            ('recreated', import_result.recreatedEntries),
            ('deleted', import_result.deletedEntries),
            ('unchanged', import_result.unchangedEntries),
        )
        if n
    )
    links = ', '.join(
        '{0} {1}'.format(n, name)
        for name, n in (
            ('created', import_result.createdEntryLinks),
            ('deleted', import_result.deletedEntryLinks),
            ('unchanged', import_result.unchangedEntryLinks),
        )
        if n
    )
    parts = []
    if entries:
      parts.append('entries: {0}'.format(entries))
    if links:
      parts.append('entry links: {0}'.format(links))
    return '; '.join(parts)

  def _GetProjectNumber(self, project_id: str) -> str:
    project_ref = projects_util.ParseProject(project_id)
    return str(projects_api.Get(project_ref).projectNumber)

  def _JobStoragePrefix(
      self, storage_uri: str, metadata_job_id: str | None
  ) -> str:
    """Returns gs://bucket/<prefix>/<job-id>/ for the per-job upload."""
    prefix = storage_uri if storage_uri.endswith('/') else storage_uri + '/'
    # A server-generated job id isn't known at upload time; use a unique folder
    # so concurrent server-id jobs sharing this storage-uri don't overwrite each
    # other's import file.
    job_folder = metadata_job_id or 'dbt-import-{0}'.format(uuid.uuid4().hex)
    return '{0}{1}/'.format(prefix, job_folder)
