#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
import json
import os
import shlex
import textwrap
from tempfile import NamedTemporaryFile
from urllib.parse import urlparse

import pytest
import requests

from airflow.providers.google.cloud.example_dags.example_dataflow_flex_template import (
    BQ_FLEX_TEMPLATE_DATASET,
    BQ_FLEX_TEMPLATE_LOCATION,
    DATAFLOW_FLEX_TEMPLATE_JOB_NAME,
    GCS_FLEX_TEMPLATE_TEMPLATE_PATH,
    PUBSUB_FLEX_TEMPLATE_SUBSCRIPTION,
    PUBSUB_FLEX_TEMPLATE_TOPIC,
)
from airflow.providers.google.cloud.example_dags.example_dataflow_sql import (
    BQ_SQL_DATASET,
    DATAFLOW_SQL_JOB_NAME,
    DATAFLOW_SQL_LOCATION,
)
from tests.providers.google.cloud.utils.gcp_authenticator import GCP_DATAFLOW_KEY, GCP_GCS_TRANSFER_KEY
from tests.test_utils.gcp_system_helpers import CLOUD_DAG_FOLDER, GoogleSystemTest, provide_gcp_context


@pytest.mark.backend("mysql", "postgres")
@pytest.mark.credential_file(GCP_DATAFLOW_KEY)
class CloudDataflowExampleDagsSystemTest(GoogleSystemTest):
    @provide_gcp_context(GCP_DATAFLOW_KEY)
    def test_run_example_gcp_dataflow_native_java(self):
        self.run_dag('example_gcp_dataflow_native_java', CLOUD_DAG_FOLDER)

    @provide_gcp_context(GCP_DATAFLOW_KEY)
    def test_run_example_gcp_dataflow_native_python(self):
        self.run_dag('example_gcp_dataflow_native_python', CLOUD_DAG_FOLDER)

    @provide_gcp_context(GCP_DATAFLOW_KEY)
    def test_run_example_gcp_dataflow_native_python_async(self):
        self.run_dag('example_gcp_dataflow_native_python_async', CLOUD_DAG_FOLDER)

    @provide_gcp_context(GCP_DATAFLOW_KEY)
    def test_run_example_gcp_dataflow_template(self):
        self.run_dag('example_gcp_dataflow_template', CLOUD_DAG_FOLDER)


GCP_PROJECT_ID = os.environ.get("GCP_PROJECT_ID", "example-project")
GCR_FLEX_TEMPLATE_IMAGE = f"gcr.io/{GCP_PROJECT_ID}/samples-dataflow-streaming-beam-sql:latest"

# https://github.com/GoogleCloudPlatform/java-docs-samples/tree/954553c/dataflow/flex-templates/streaming_beam_sql
GCS_TEMPLATE_PARTS = urlparse(GCS_FLEX_TEMPLATE_TEMPLATE_PATH)
GCS_FLEX_TEMPLATE_BUCKET_NAME = GCS_TEMPLATE_PARTS.netloc


EXAMPLE_FLEX_TEMPLATE_REPO = "GoogleCloudPlatform/java-docs-samples"
EXAMPLE_FLEX_TEMPLATE_COMMIT = "deb0745be1d1ac1d133e1f0a7faa9413dbfbe5fe"
EXAMPLE_FLEX_TEMPLATE_SUBDIR = "dataflow/flex-templates/streaming_beam_sql"


@pytest.mark.backend("mysql", "postgres")
@pytest.mark.credential_file(GCP_GCS_TRANSFER_KEY)
class CloudDataflowExampleDagFlexTemplateJavagSystemTest(GoogleSystemTest):
    @provide_gcp_context(GCP_GCS_TRANSFER_KEY, project_id=GoogleSystemTest._project_id())
    def setUp(self) -> None:
        super().setUp()
        # Create a Cloud Storage bucket
        self.execute_cmd(["gsutil", "mb", f"gs://{GCS_FLEX_TEMPLATE_BUCKET_NAME}"])

        # Build image with pipeline
        with NamedTemporaryFile("w") as f:
            cloud_build_config = {
                'steps': [
                    {'name': 'gcr.io/cloud-builders/git', 'args': ['clone', "$_EXAMPLE_REPO", "repo_dir"]},
                    {
                        'name': 'gcr.io/cloud-builders/git',
                        'args': ['checkout', '$_EXAMPLE_COMMIT'],
                        'dir': 'repo_dir',
                    },
                    {
                        'name': 'maven',
                        'args': ['mvn', 'clean', 'package'],
                        'dir': 'repo_dir/$_EXAMPLE_SUBDIR',
                    },