"""Real-environment validation DAG for PR #66613 (GH-39002)."""
from __future__ import annotations

import os
from datetime import datetime

from airflow.providers.databricks.hooks.databricks import DatabricksHook
from airflow.providers.databricks.operators.databricks import (
    DatabricksCreateJobsOperator,
    DatabricksRunNowOperator,
    DatabricksSubmitRunOperator,
)
from airflow.sdk import DAG, task

NOTEBOOK_PATH = os.environ.get(
    "PR66613_NOTEBOOK_PATH", "/Users/<your-user>@example.com/airflow-pr66613-noop"
)


def _hook() -> DatabricksHook:
    return DatabricksHook(databricks_conn_id="databricks_default")


with DAG(
    dag_id="pr66613_realenv",
    start_date=datetime(2026, 1, 1),
    schedule=None,
    catchup=False,
    params={"env": "prod", "batch_size": "100"},
    tags=["databricks", "pr66613"],
) as dag:
    create_job = DatabricksCreateJobsOperator(
        task_id="create_job",
        json={
            "name": "{{ dag.dag_id }}-{{ ts_nodash }}",
            "tasks": [{"task_key": "noop", "notebook_task": {"notebook_path": NOTEBOOK_PATH}}],
        },
    )

    @task(task_id="verify_create_job")
    def verify_create_job(job_id: int) -> int:
        job = _hook()._do_api_call(("GET", "2.2/jobs/get"), {"job_id": job_id})
        params = job["settings"].get("parameters", [])
        assert {"name": "env", "default": "prod"} in params, params
        assert {"name": "batch_size", "default": "100"} in params, params
        return job_id

    run_now = DatabricksRunNowOperator(
        task_id="run_now",
        job_id="{{ ti.xcom_pull(task_ids='verify_create_job') }}",
        params={"env": "staging", "batch_size": "42"},
        wait_for_termination=False,
    )

    @task(task_id="verify_run_now")
    def verify_run_now(databricks_run_id: int) -> None:
        run = _hook()._do_api_call(("GET", "2.2/jobs/runs/get"), {"run_id": databricks_run_id})
        found = {p["name"]: p.get("value", p.get("default")) for p in run.get("job_parameters", [])}
        assert found.get("env") == "staging", found
        assert found.get("batch_size") == "42", found
        _hook().cancel_run(databricks_run_id)

    submit_run = DatabricksSubmitRunOperator(
        task_id="submit_run",
        notebook_task={"notebook_path": NOTEBOOK_PATH},
        new_cluster={
            "spark_version": "15.4.x-scala2.12",
            "node_type_id": "i3.xlarge",
            "num_workers": 0,
            "spark_conf": {"spark.master": "local[*]"},
            "custom_tags": {"ResourceClass": "SingleNode"},
        },
        params={"env": "dev", "shard": "1"},
        wait_for_termination=False,
    )

    @task(task_id="verify_submit_run")
    def verify_submit_run(databricks_run_id: int) -> None:
        run = _hook()._do_api_call(("GET", "2.2/jobs/runs/get"), {"run_id": databricks_run_id})
        base_params = run["tasks"][0]["notebook_task"]["base_parameters"]
        # Operator-level params override DAG-level params for shared keys; DAG-level
        # keys not overridden are still inherited (here: batch_size from the DAG).
        assert base_params == {"env": "dev", "batch_size": "100", "shard": "1"}, base_params
        _hook().cancel_run(databricks_run_id)

    @task(task_id="cleanup_job", trigger_rule="all_done")
    def cleanup_job(job_id: int) -> None:
        try:
            _hook()._do_api_call(("POST", "2.2/jobs/delete"), {"job_id": job_id})
        except Exception:  # noqa: BLE001
            pass

    job_id_xcom = verify_create_job(create_job.output)
    run_now_done = verify_run_now(
        databricks_run_id="{{ ti.xcom_pull(task_ids='run_now', key='run_id') }}"
    )
    submit_run_done = verify_submit_run(
        databricks_run_id="{{ ti.xcom_pull(task_ids='submit_run', key='run_id') }}"
    )
    create_job >> job_id_xcom >> run_now >> run_now_done
    submit_run >> submit_run_done
    [run_now_done, submit_run_done] >> cleanup_job(job_id_xcom)