# 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.
from __future__ import annotations

from datetime import datetime

from airflow.providers.amazon.aws.operators.sagemaker_unified_studio import (
    SageMakerNotebookOperator,
)

from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS

if AIRFLOW_V_3_0_PLUS:
    from airflow.sdk import DAG, chain, task
else:
    # Airflow 2 path
    from airflow.decorators import task  # type: ignore[attr-defined,no-redef]
    from airflow.models.baseoperator import chain  # type: ignore[attr-defined,no-redef]
    from airflow.models.dag import DAG  # type: ignore[attr-defined,no-redef,assignment]

from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder

"""
Prerequisites: The account which runs this test must manually have the following:
1. An IAM IDC organization set up in the testing region with a user initialized
2. A SageMaker Unified Studio Domain (with default VPC and roles)
3. A project within the SageMaker Unified Studio Domain
4. A notebook (test_notebook.ipynb) placed in the project's s3 path

The test runs in two stages:
1. run-notebook-explicit: passes domain_id, domain_region, and project_id directly as operator
   parameters. No environment variables are required. Requires sagemaker-studio>=1.0.25.
2. setup_mwaa_environment + run-notebook: sets MWAA-style environment variables and runs the
   notebook using the legacy env-var-based resolution path.

The ordering is intentional: run-notebook-explicit runs BEFORE the env vars are set, so on
older SDK versions (<1.0.25) that cannot resolve the region from explicit params, the test
will fail at stage 1 rather than accidentally passing via the env vars.
"""

DAG_ID = "example_sagemaker_unified_studio"

# Externally fetched variables:
DOMAIN_ID_KEY = "DOMAIN_ID"
PROJECT_ID_KEY = "PROJECT_ID"
ENVIRONMENT_ID_KEY = "ENVIRONMENT_ID"
S3_PATH_KEY = "S3_PATH"
REGION_NAME_KEY = "REGION_NAME"

sys_test_context_task = (
    SystemTestContextBuilder()
    .add_variable(DOMAIN_ID_KEY)
    .add_variable(PROJECT_ID_KEY)
    .add_variable(ENVIRONMENT_ID_KEY)
    .add_variable(S3_PATH_KEY)
    .add_variable(REGION_NAME_KEY)
    .build()
)


def get_mwaa_environment_params(
    domain_id: str,
    project_id: str,
    environment_id: str,
    s3_path: str,
    region_name: str,
):
    AIRFLOW_PREFIX = "AIRFLOW__WORKFLOWS__"

    parameters = {}
    parameters[f"{AIRFLOW_PREFIX}DATAZONE_DOMAIN_ID"] = domain_id
    parameters[f"{AIRFLOW_PREFIX}DATAZONE_PROJECT_ID"] = project_id
    parameters[f"{AIRFLOW_PREFIX}DATAZONE_ENVIRONMENT_ID"] = environment_id
    parameters[f"{AIRFLOW_PREFIX}DATAZONE_SCOPE_NAME"] = "dev"
    parameters[f"{AIRFLOW_PREFIX}DATAZONE_STAGE"] = "prod"
    parameters[f"{AIRFLOW_PREFIX}DATAZONE_ENDPOINT"] = f"https://datazone.{region_name}.api.aws"
    parameters[f"{AIRFLOW_PREFIX}PROJECT_S3_PATH"] = s3_path
    parameters[f"{AIRFLOW_PREFIX}DATAZONE_DOMAIN_REGION"] = region_name
    return parameters


@task
def mock_mwaa_environment(parameters: dict):
    """
    Sets several environment variables in the container to emulate an MWAA environment provisioned
    within SageMaker Unified Studio. When running in the ECSExecutor, this is a no-op.
    """
    import os

    for key, value in parameters.items():
        os.environ[key] = value


with DAG(
    DAG_ID,
    schedule="@once",
    start_date=datetime(2021, 1, 1),
    tags=["example"],
    catchup=False,
) as dag:
    test_context = sys_test_context_task()

    test_env_id = test_context[ENV_ID_KEY]
    domain_id = test_context[DOMAIN_ID_KEY]
    project_id = test_context[PROJECT_ID_KEY]
    environment_id = test_context[ENVIRONMENT_ID_KEY]
    s3_path = test_context[S3_PATH_KEY]
    region_name = test_context[REGION_NAME_KEY]

    mock_mwaa_environment_params = get_mwaa_environment_params(
        domain_id,
        project_id,
        environment_id,
        s3_path,
        region_name,
    )

    setup_mwaa_environment = mock_mwaa_environment(mock_mwaa_environment_params)

    notebook_path = "test_notebook.ipynb"  # This should be the path to your .ipynb, .sqlnb, or .vetl file in your project.

    # [START howto_operator_sagemaker_unified_studio_notebook_explicit_params]
    # Run notebook with domain_id/project_id/domain_region passed explicitly as operator parameters.
    # No environment variables needed — the SDK resolves the S3 path and region from these params.
    # Requires sagemaker-studio>=1.0.25.
    # NOTE: this task runs BEFORE env vars are set intentionally, to prove that explicit params
    # work without any MWAA-style environment variables present.
    run_notebook_explicit_params = SageMakerNotebookOperator(
        task_id="run-notebook-explicit",
        domain_id=domain_id,
        project_id=project_id,
        domain_region=region_name,
        input_config={"input_path": notebook_path, "input_params": {}},
        output_config={"output_formats": ["NOTEBOOK"]},  # optional
        compute={
            "instance_type": "ml.m5.large",
            "volume_size_in_gb": 30,
        },  # optional
        termination_condition={"max_runtime_in_seconds": 600},  # optional
        tags={},  # optional
        wait_for_completion=True,  # optional
        waiter_delay=5,  # optional
        deferrable=False,  # optional
    )
    # [END howto_operator_sagemaker_unified_studio_notebook_explicit_params]

    # [START howto_operator_sagemaker_unified_studio_notebook]
    # Run notebook using the legacy env-var-based resolution path (MWAA-style).
    run_notebook = SageMakerNotebookOperator(
        task_id="run-notebook",
        input_config={"input_path": notebook_path, "input_params": {}},
        output_config={"output_formats": ["NOTEBOOK"]},  # optional
        compute={
            "instance_type": "ml.m5.large",
            "volume_size_in_gb": 30,
        },  # optional
        termination_condition={"max_runtime_in_seconds": 600},  # optional
        tags={},  # optional
        wait_for_completion=True,  # optional
        waiter_delay=5,  # optional
        deferrable=False,  # optional
        executor_config={  # optional
            "overrides": {
                "containerOverrides": [
                    {
                        "environment": [
                            {"name": key, "value": value}
                            for key, value in mock_mwaa_environment_params.items()
                        ],
                        "name": "ECSExecutorContainer",  # Necessary parameter
                    }
                ]
            }
        },
    )
    # [END howto_operator_sagemaker_unified_studio_notebook]

    chain(
        # TEST SETUP
        test_context,
        # TEST BODY: explicit params first (no env vars set yet)
        run_notebook_explicit_params,
        # TEST SETUP: set MWAA-style env vars
        setup_mwaa_environment,
        # TEST BODY: legacy env-var-based resolution
        run_notebook,
    )

    from tests_common.test_utils.watcher import watcher

    # This test needs watcher in order to properly mark success/failure
    # when "tearDown" task with trigger rule is part of the DAG
    list(dag.tasks) >> watcher()

from tests_common.test_utils.system_tests import get_test_run  # noqa: E402

# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst)
test_run = get_test_run(dag)
