import os
import re
from collections.abc import Iterator
from contextlib import ExitStack, contextmanager
from typing import Any

import dagster._check as check
import pytest
from dagster import AssetExecutionContext, asset, materialize
from dagster._core.errors import DagsterPipesExecutionError
from dagster_databricks._test_utils import (
    databricks_client,  # noqa: F401
    temp_dbfs_script,
    temp_workspace_notebook,
    upload_dagster_pipes_whl,
)
from dagster_databricks.pipes import PipesDatabricksClient, PipesDbfsMessageReader
from databricks.sdk import WorkspaceClient
from databricks.sdk.service import compute, jobs

IS_BUILDKITE = os.getenv("BUILDKITE") is not None
IS_WORKSPACE = os.getenv("DATABRICKS_HOST") is not None


def script_fn():
    import os
    import sys

    from dagster_pipes import (
        DAGSTER_PIPES_CONTEXT_ENV_VAR,
        PipesCliArgsParamsLoader,
        PipesDatabricksNotebookWidgetsParamsLoader,
        PipesDbfsContextLoader,
        PipesDbfsMessageWriter,
        PipesEnvVarParamsLoader,
        open_dagster_pipes,
    )

    # To facilitate using the same script for testing in both the new cluster and existing cluster
    # instances and python_script and notebook_task execution we dynamically configure the
    # PipesParamsLoader here by checking for the presence of pipes-specific env vars. If these
    # are set, we know we are in the new cluster case and load params via the env vars.
    def is_notebook():
        try:
            from pyspark.dbutils import DBUtils  # pyright: ignore[reportMissingImports]

            dbutils = DBUtils(spark)  # pyright: ignore[reportUndefinedVariable] # noqa: F821
            dbutils.widgets.get(DAGSTER_PIPES_CONTEXT_ENV_VAR)
            return True
        except Exception:
            return False

    def get_params_loader():
        if is_notebook():
            return PipesDatabricksNotebookWidgetsParamsLoader(dbutils.widgets)  # pyright: ignore[reportUndefinedVariable] # noqa: F821
        elif DAGSTER_PIPES_CONTEXT_ENV_VAR in os.environ:
            return PipesEnvVarParamsLoader()
        else:
            return PipesCliArgsParamsLoader()

    with open_dagster_pipes(
        params_loader=get_params_loader(),
        context_loader=PipesDbfsContextLoader(),
        message_writer=PipesDbfsMessageWriter(),
    ) as context:
        multiplier = context.get_extra("multiplier")
        value = 2 * multiplier
        print("hello from databricks stdout")  # noqa: T201
        print("hello from databricks stderr", file=sys.stderr)  # noqa: T201
        context.log.info(f"{context.asset_key}: {2} * {multiplier} = {value}")
        context.report_asset_materialization(
            metadata={"value": value},
        )


CLUSTER_DEFAULTS = {
    "spark_version": "12.2.x-scala2.12",
    "num_workers": 0,
}

TASK_KEY = "DAGSTER_PIPES_TASK"


# Create a new cluster on Databricks, yield the cluster_id, and then terminate the cluster
# on exit. This is useful for testing the PipesClient against an existing cluster.
@contextmanager
def temp_databricks_cluster(client: WorkspaceClient, forward_logs: bool) -> Iterator[str]:
    cluster_id = None
    try:
        # Need to make sure that nested dicts are wrapped in their SDK objects here (see
        # `use_inner_objects`). It's required by this API but apparently not by others. Hopefully
        # updating to a newer SDK will alleviate some of the annoyances around mixed acceptance of
        # objects from the SDK and their dict forms.
        cluster_spec = make_new_cluster_spec(forward_logs, use_inner_objects=True)
        cluster_details = client.clusters.create_and_wait(**cluster_spec)
        cluster_id = check.not_none(cluster_details.cluster_id)
        yield cluster_id
    finally:
        if cluster_id:
            client.clusters.delete(cluster_id)


def make_submit_task_dict(
    file_path: str,
    dagster_pipes_whl_path: str,
    forward_logs: bool,
    task_type: str,
    file_path_key: str,
    cluster_id: str | None = None,
) -> dict[str, Any]:
    submit_spec = {
        "libraries": [
            {"whl": dagster_pipes_whl_path},
        ],
        "task_key": TASK_KEY,
        task_type: {
            file_path_key: file_path,
            "source": jobs.Source.WORKSPACE,
        },
    }
    if cluster_id:
        submit_spec["existing_cluster_id"] = cluster_id
    else:
        submit_spec["new_cluster"] = make_new_cluster_spec(forward_logs)
    return submit_spec


def make_new_cluster_spec(forward_logs: bool, use_inner_objects: bool = False) -> Any:
    cluster_spec = CLUSTER_DEFAULTS.copy()
    databricks_host = os.getenv("DATABRICKS_HOST", "")
    if "azuredatabricks.net" in databricks_host:
        cluster_spec["node_type_id"] = "Standard_DS3_v2"
    else:  # Assume AWS
        cluster_spec["node_type_id"] = "i3.xlarge"
    if forward_logs:
        log_conf = {"dbfs": {"destination": "dbfs:/cluster-logs"}}
        cluster_spec["cluster_log_conf"] = (
            compute.ClusterLogConf.from_dict(log_conf) if use_inner_objects else log_conf
        )
    return cluster_spec


def make_submit_task(
    file_path: str,
    dagster_pipes_whl_path: str,
    forward_logs: bool,
    task_type: str,
    file_path_key: str,
    cluster_id: str | None = None,
) -> jobs.SubmitTask:
    return jobs.SubmitTask.from_dict(
        make_submit_task_dict(
            file_path=file_path,
            dagster_pipes_whl_path=dagster_pipes_whl_path,
            forward_logs=forward_logs,
            task_type=task_type,
            file_path_key=file_path_key,
            cluster_id=cluster_id,
        )
    )


# Test both with and without log forwarding. This is important because the PipesClient spins up log
# readers before it knows the task specification
@pytest.mark.skipif(IS_BUILDKITE, reason="Not configured to run on BK yet.")
@pytest.mark.skipif(not IS_WORKSPACE, reason="No DB workspace credentials found.")
@pytest.mark.parametrize("forward_logs", [True, False])
@pytest.mark.parametrize("use_existing_cluster", [True, False])
@pytest.mark.parametrize(
    "task_type, file_path_key",
    [("spark_python_task", "python_file"), ("notebook_task", "notebook_path")],
)
def test_pipes_client(
    capsys,
    databricks_client: WorkspaceClient,  # noqa: F811
    forward_logs: bool,
    use_existing_cluster: bool,
    task_type: str,
    file_path_key: str,
):
    if use_existing_cluster and forward_logs:
        # Log forwarding from existing clusters requires using Pipes messages as log transport. This has to be done because:
        # (a) logs are flushed every 5 minutes or on cluster termination.
        # If cluster termination is not tied to the end of the launched job, there is no mechanism
        # to wait for logs to be flushed. (b) The logs reader functions by forwarding stdout/stderr,
        # which is shared across all jobs on the cluster-- so there is no way to scope the
        # forwarding to just the target job.
        #
        message_reader = PipesDbfsMessageReader(
            client=databricks_client, include_stdio_in_messages=True
        )
    else:
        message_reader = None

    @asset
    def number_x(context: AssetExecutionContext, pipes_client: PipesDatabricksClient):
        with ExitStack() as stack:
            dagster_pipes_whl_path = stack.enter_context(
                upload_dagster_pipes_whl(databricks_client)
            )

            if task_type == "spark_python_task":
                file_path = stack.enter_context(
                    temp_dbfs_script(databricks_client, script_fn=script_fn)
                )
                file_path = f"dbfs:{file_path}"
            else:
                file_path = stack.enter_context(
                    temp_workspace_notebook(
                        databricks_client,
                        workspace_path="/Shared/dagster-pipes-test",
                        script_fn=script_fn,
                    )
                )
            if use_existing_cluster:
                cluster_id = stack.enter_context(
                    temp_databricks_cluster(databricks_client, forward_logs)
                )
                task = make_submit_task(
                    file_path,
                    dagster_pipes_whl_path,
                    forward_logs,
                    task_type,
                    file_path_key,
                    cluster_id,
                )
            else:
                task = make_submit_task(
                    file_path, dagster_pipes_whl_path, forward_logs, task_type, file_path_key
                )
            return pipes_client.run(
                task=task,
                context=context,
                extras={"multiplier": 2, "storage_root": "fake"},
            ).get_results()

    result = materialize(
        [number_x],
        resources={
            "pipes_client": PipesDatabricksClient(databricks_client, message_reader=message_reader)
        },
        raise_on_error=False,
    )
    assert result.success
    mats = result.asset_materializations_for_node(number_x.op.name)
    assert mats[0].metadata["value"].value == 4
    if forward_logs:
        captured = capsys.readouterr()
        assert re.search(r"hello from databricks stdout\n", captured.out, re.MULTILINE)
        assert re.search(r"hello from databricks stderr\n", captured.err, re.MULTILINE)

    # check Databricks metadata automatically added to materialization
    assert "Databricks Job Run ID" in mats[0].metadata
    assert "Databricks Job Run URL" in mats[0].metadata


@pytest.mark.skipif(IS_BUILDKITE, reason="Not configured to run on BK yet.")
@pytest.mark.skipif(not IS_WORKSPACE, reason="No DB workspace credentials found.")
@pytest.mark.parametrize("task_type, file_path_key", [("spark_python_task", "python_file")])
def test_nonexistent_entry_point(
    databricks_client: WorkspaceClient,  # noqa: F811
    task_type: str,
    file_path_key: str,
):
    @asset
    def fake(context: AssetExecutionContext, pipes_client: PipesDatabricksClient):
        task = make_submit_task(
            "/fake/fake",
            "/fake/fake",
            forward_logs=False,
            task_type=task_type,
            file_path_key=file_path_key,
        )
        return pipes_client.run(task=task, context=context).get_results()

    with pytest.raises(DagsterPipesExecutionError, match=r"Cannot read the python file"):
        materialize(
            [fake],
            resources={"pipes_client": PipesDatabricksClient(databricks_client)},
        )
