# -*- 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.
"""Command to unpause a pipeline."""

import datetime
import random
import subprocess
import time
from googlecloudsdk.api_lib.composer import dags_util
from googlecloudsdk.api_lib.composer import environments_util
from googlecloudsdk.api_lib.composer import util
from googlecloudsdk.calliope import base as calliope_base
from googlecloudsdk.calliope import exceptions as calliope_exceptions
from googlecloudsdk.command_lib.orchestration_pipelines.tools import composer_utils
from googlecloudsdk.command_lib.orchestration_pipelines.tools import gcs_utils
from googlecloudsdk.command_lib.orchestration_pipelines.tools import yaml_processor
from googlecloudsdk.core import log
from googlecloudsdk.core import resources

DEFAULT_POLLING_TIME_SECONDS = datetime.timedelta(seconds=2)
POLL_JITTER_SECONDS = datetime.timedelta(seconds=0.5)


@calliope_base.Hidden
@calliope_base.DefaultUniverseOnly
@calliope_base.ReleaseTracks(calliope_base.ReleaseTrack.BETA)
class Unpause(calliope_base.Command):
  """A command that unpauses an orchestration pipeline."""

  def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self._subprocess = subprocess

  @staticmethod
  def Args(parser):
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument(
        "--environment",
        help=(
            "The target environment of the pipeline, as defined in"
            " deployment.yaml."
        ),
    )
    group.add_argument(
        "--runner",
        help="The full resource name to unpause a pipeline from.",
    )
    parser.add_argument(
        "--bundle",
        required=True,
        help="The ID of the bundle the pipeline belongs to.",
    )
    parser.add_argument(
        "--pipeline",
        required=True,
        help="The ID of the pipeline to unpause.",
    )
    parser.add_argument(
        "--async",
        action="store_true",
        dest="async_",
        help=(
            "If set, the command will return after updating the pipeline files"
            " and unpause the pipeline asynchronously in the"
            " background."
        ),
    )
    yaml_processor.add_substitution_flags(parser)

  def Run(self, args):
    api_version = util.GetApiVersion(self.ReleaseTrack())

    # 1. Retrieve the pipeline to be activated.
    list_filter = composer_utils.build_dags_filter_tags(
        bundle=args.bundle, pipeline=args.pipeline, is_current=True
    )
    if args.runner:
      resource_name = composer_utils.build_resource_name(runner=args.runner)
    else:
      environment_model = yaml_processor.load_environment_with_args(args)
      resource_name = composer_utils.build_resource_name(
          env_model=environment_model
      )
    environment_ref = resources.REGISTRY.ParseRelativeName(
        resource_name,
        collection="composer.projects.locations.environments",
        api_version=api_version,
    )
    list_dags_response = dags_util.ListDags(
        environment_ref,
        list_filter=list_filter,
    )

    if len(list_dags_response.dags) == 0:
      return {
          "result": "failed",
          "reason": "No pipeline found for given bundle and pipeline IDs.",
      }

    if len(list_dags_response.dags) > 1:
      return {
          "result": "failed",
          "reason": (
              "Multiple pipelines found for given bundle and pipeline IDs."
          ),
      }

    # If the pipeline is not paused and its DAG state is ACTIVE, return early.
    if (
        not composer_utils.get_pipeline_paused_status(
            list_dags_response.dags[0]
        )
        and list_dags_response.dags[0].state.name == "ACTIVE"
    ):
      log.status.Print(f"Pipeline {args.pipeline} is already active.")
      return {"result": "success"}

    # 2. Retrieve the environment storage bucket and build the manifest GCS
    # path.
    environment = environments_util.Get(environment_ref)
    if not environment.storageConfig or not environment.storageConfig.bucket:
      return {
          "result": "failed",
          "reason": (
              "Could not retrieve GCS storage bucket for Composer environment."
          ),
      }
    manifest_gcs_path = "gs://{}/{}/{}/{}".format(
        environment.storageConfig.bucket,
        gcs_utils.ORCHESTRATION_PIPELINES_DATA_DIRECTORY,
        args.bundle,
        gcs_utils.MANIFEST_FILE_NAME,
    )

    try:
      # 3. Update the manifest file in GCS with retry logic and optimistic
      # locking.
      gcs_utils.UpdatePausedPipelinesInManifestWithRetry(
          self._subprocess, manifest_gcs_path, args.pipeline, "remove"
      )
      log.status.Print(f"Successfully updated manifest for {args.bundle}.")

      # 4. Touch the pipeline's python file to trigger the Airflow scheduler to
      # pick up the changes in the next run.
      pipeline_gcs_path = "gs://{}/{}/{}/{}.py".format(
          environment.storageConfig.bucket,
          gcs_utils.ORCHESTRATION_PIPELINES_DAGS_DIRECTORY,
          args.bundle,
          args.pipeline,
      )
      log.status.Print("Updating metadata for pipeline files...")
      gcs_utils.TouchPipelineFile(self._subprocess, pipeline_gcs_path)
    except (
        calliope_exceptions.BadFileException,
        calliope_exceptions.HttpException,
    ) as e:
      log.status.Print("Reverting manifest file...")
      gcs_utils.UpdatePausedPipelinesInManifestWithRetry(
          self._subprocess, manifest_gcs_path, args.pipeline, "add"
      )
      return {
          "result": "failed",
          "reason": f"{str(e)}",
      }
    except ValueError as _:
      return {
          "result": "failed",
          "reason": (
              "An unexpected error occurred while unpausing the pipeline."
          ),
      }

    dag_ref = resources.REGISTRY.ParseRelativeName(
        list_dags_response.dags[0].name,
        collection="composer.projects.locations.environments.dags",
        api_version=api_version,
    )
    dags_util.ActivateDag(dag_ref)
    log.status.Print(f"Activating pipeline {args.pipeline}...")

    if args.async_:
      log.status.Print(
          "Pipeline files are updated and pipeline is being activated. Running"
          " in async mode, exiting..."
      )
      return {"result": "success"}

    # 5. Wait for the scheduler to pick up the changes and update the DAG.
    log.status.Print("Waiting for the pipeline to be unpaused...")
    while True:
      try:
        unpaused_dags_list = dags_util.ListDags(
            environment_ref,
            list_filter=list_filter,
        )

        # Check for both pipeline paused status and DAG state to ensure the
        # pipeline is active.
        if (
            unpaused_dags_list.dags
            and not composer_utils.get_pipeline_paused_status(
                unpaused_dags_list.dags[0]
            )
            and unpaused_dags_list.dags[0].state.name == "ACTIVE"
        ):
          return {"result": "success"}
        time.sleep(
            DEFAULT_POLLING_TIME_SECONDS.total_seconds()
            + random.uniform(
                -POLL_JITTER_SECONDS.total_seconds(),
                POLL_JITTER_SECONDS.total_seconds(),
            )
        )
      except KeyboardInterrupt:
        break
    return {
        "result": "success",
        "reason": (
            "The command was interrupted, but the process is still running"
            " in the background."
        ),
    }
