import os
from datetime import datetime
from airflow import DAG
from airflow.models.connection import Connection
from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
from airflow.providers.standard.operators.bash import BashOperator

def _set_k8s_conn_from_env() -> None:
    extra = {
        "in_cluster": True,
        "namespace": "airflow",
    }

    conn = Connection(
        conn_id="k8s-conn-id",
        conn_type="kubernetes",
        description="K8s connection defined only via AIRFLOW_CONN_* env var",
        extra=extra,
    )

    os.environ[f"AIRFLOW_CONN_{conn.conn_id.upper()}"] = conn.get_uri()


_set_k8s_conn_from_env()


with DAG(
    dag_id="k8s_env_deferrable_repro_41706",
    start_date=datetime(2024, 1, 1),
    schedule=None,
    catchup=False,
    tags=["k8s", "deferrable", "env-conn", "issue-41706"],
) as dag:
    k8s_pod = KubernetesPodOperator(
        task_id="k8s_pod_env_conn_deferrable",
        kubernetes_conn_id="k8s-conn-id",
        namespace="airflow",
        name="k8s-env-deferrable-pod",
        image="alpine",
        cmds=["sh", "-c"],
        arguments=["mkdir -p /airflow/xcom/; echo '[1,2,3,4]' > /airflow/xcom/return.json"],
        do_xcom_push=True,
        get_logs=True,
        deferrable=True,
        on_finish_action="keep_pod",
        startup_timeout_seconds=600,
    )

    print_xcom = BashOperator(
        task_id="print_xcom_result",
        bash_command="echo \"XCOM from KPO: {{ ti.xcom_pull('k8s_pod_env_conn_deferrable')[0] }}\"",
    )

    k8s_pod >> print_xcom