from unittest.mock import patch

from airflow_pydantic import Dag
from airflow_pydantic.airflow import Pool
from airflow_pydantic.testing import pools

try:
    from airflow import DAG  # noqa: F401
    from airflow.models.pool import Pool

    _HAVE_AIRFLOW = True
except ImportError:
    _HAVE_AIRFLOW = False

try:
    from airflow.providers.ssh.operators.ssh import SSHOperator  # noqa: F401

    _HAVE_AIRFLOW_SSH = True
except ImportError:
    _HAVE_AIRFLOW_SSH = False


class TestRender:
    def test_render_operator_ssh(self, ssh_operator):
        imports, globals_, task = ssh_operator.render()
        assert imports == [
            "from airflow.providers.ssh.operators.ssh import SSHOperator",
            "from airflow.models.pool import Pool",
            "from airflow.providers.ssh.hooks.ssh import SSHHook",
        ]
        assert globals_ == []
        assert (
            task
            == "SSHOperator(pool=Pool.get_pool('blerg').pool, do_xcom_push=True, ssh_hook=SSHHook(remote_host='test', username='test'), ssh_conn_id='test', command=\"bash -lc 'set -ex\\ntest1\\ntest2'\", cmd_timeout=10, environment={'test': 'test'}, get_pty=True, task_id='test-ssh-operator')"
        )

    def test_render_operator_ssh_host_variable(self, ssh_operator_balancer):
        imports, globals_, task = ssh_operator_balancer.render()
        assert imports == [
            "from airflow.providers.ssh.operators.ssh import SSHOperator",
            "from airflow.models.pool import Pool",
            "from airflow.providers.ssh.hooks.ssh import SSHHook",
            "from airflow.models.variable import Variable as AirflowVariable",
        ]
        assert globals_ == []
        assert (
            task
            == "SSHOperator(pool=Pool.create_or_update_pool(name='test_host', slots=8, description='Balancer pool for host(test_host)', include_deferred=False).pool, do_xcom_push=True, ssh_hook=SSHHook(remote_host='test_host.local', username='test_user', password=AirflowVariable.get('VAR', deserialize_json=True)['password']), ssh_conn_id='test', command='bash -lc \\'export var=\"{{ ti.blerg }}\"\\ncd /tmp\\nset -ex\\ntest1\\ntest2\\'', cmd_timeout=10, environment={'test': 'test'}, get_pty=True, task_id='test-ssh-operator')"
        )

    def test_render_operator_ssh_host_variable_from_template(self, ssh_operator_balancer_template):
        imports, globals_, task = ssh_operator_balancer_template.render()
        assert imports == [
            "from airflow.providers.ssh.operators.ssh import SSHOperator",
            "from airflow.models.pool import Pool",
            "from airflow.providers.ssh.hooks.ssh import SSHHook",
            "from airflow.models.variable import Variable as AirflowVariable",
        ]
        assert globals_ == []
        assert (
            task
            == "SSHOperator(pool=Pool.create_or_update_pool(name='test_host', slots=8, description='Balancer pool for host(test_host)', include_deferred=False).pool, do_xcom_push=True, ssh_hook=SSHHook(remote_host='test_host.local', username='test_user', password=AirflowVariable.get('VAR', deserialize_json=True)['password']), ssh_conn_id='test', command='bash -lc \\'export var=\"{{ ti.blerg }}\"\\ncd /tmp\\nset -ex\\ntest1\\ntest2\\'', cmd_timeout=10, environment={'test': 'test'}, get_pty=True, task_id='test-ssh-operator')"
        )

    def test_render_dag(self, dag):
        assert isinstance(dag, Dag)
        assert (
            dag.render(debug_filename="airflow_pydantic/tests/core/rendered/dag.py")
            == """# Generated by airflow-config
from datetime import datetime, timedelta

from airflow.models import DAG
from airflow.models.pool import Pool
from airflow.providers.ssh.hooks.ssh import SSHHook
from airflow.providers.ssh.operators.ssh import SSHOperator
from airflow.providers.standard.operators.bash import BashOperator
from airflow.providers.standard.operators.python import PythonOperator
from airflow.providers.standard.sensors.bash import BashSensor
from airflow.providers.standard.sensors.python import PythonSensor

from airflow_pydantic.tests.conftest import foo

with DAG(
    description="",
    schedule="* * * * *",
    start_date=datetime.fromisoformat("2025-01-01T00:00:00"),
    end_date=datetime.fromisoformat("2026-01-01T00:00:00"),
    max_active_tasks=1,
    max_active_runs=1,
    catchup=False,
    is_paused_upon_creation=True,
    tags=["a", "b"],
    dag_display_name="test",
    dag_id="a-dag",
    default_args={
        "owner": "airflow",
        "email": ["test@test.com"],
        "email_on_failure": True,
        "email_on_retry": True,
        "retries": 3,
        "retry_delay": timedelta(seconds=300.0),
        "start_date": datetime.fromisoformat("2025-01-01T00:00:00"),
        "end_date": datetime.fromisoformat("2026-01-01T00:00:00"),
        "depends_on_past": True,
        "queue": "default",
        "pool": "default",
        "pool_slots": 1,
        "do_xcom_push": True,
        "task_display_name": "test",
    },
) as dag:
    task1 = PythonOperator(
        pool=Pool.get_pool("test-pool-model").pool,
        python_callable=foo,
        op_args=["test"],
        op_kwargs={"test": "test"},
        templates_dict={"test": "test"},
        templates_exts=[".sql", ".hql"],
        show_return_value_in_logs=True,
        task_id="test-python-operator",
        dag=dag,
    )
    task2 = BashOperator(
        pool=Pool.create_or_update_pool(name="test", slots=5, description="", include_deferred=False).pool,
        bash_command="test",
        env={"test": "test"},
        append_env=True,
        output_encoding="utf-8",
        skip_on_exit_code=99,
        cwd="test",
        output_processor=foo,
        task_id="test-bash-operator",
        dag=dag,
    )
    task3 = SSHOperator(
        pool=Pool.get_pool("blerg").pool,
        do_xcom_push=True,
        ssh_hook=SSHHook(remote_host="test", username="test"),
        ssh_conn_id="test",
        command="bash -lc 'set -ex\\ntest1\\ntest2'",
        cmd_timeout=10,
        environment={"test": "test"},
        get_pty=True,
        task_id="test-ssh-operator",
        dag=dag,
    )
    task4 = BashSensor(bash_command="test", task_id="test-bash-sensor", dag=dag)
    task5 = PythonSensor(python_callable=foo, op_args=["test"], op_kwargs={"test": "test"}, task_id="test-python-sensor", dag=dag)
"""
        )
        if _HAVE_AIRFLOW_SSH:
            with pools(Pool()), patch("airflow.models.pool.Pool.get_pool") as mock_get_pool:
                mock_get_pool.return_value = Pool()
                dag.instantiate()
                exec(dag.render())

    def test_render_with_dependencies(self, dag):
        dag.tasks["task1"].dependencies = []
        dag.tasks["task2"].dependencies = ["task1"]
        dag.tasks["task3"].dependencies = ["task1", "task2"]
        assert isinstance(dag, Dag)
        assert (
            dag.render(debug_filename="airflow_pydantic/tests/core/rendered/dependencies.py")
            == """# Generated by airflow-config
from datetime import datetime, timedelta

from airflow.models import DAG
from airflow.models.pool import Pool
from airflow.providers.ssh.hooks.ssh import SSHHook
from airflow.providers.ssh.operators.ssh import SSHOperator
from airflow.providers.standard.operators.bash import BashOperator
from airflow.providers.standard.operators.python import PythonOperator
from airflow.providers.standard.sensors.bash import BashSensor
from airflow.providers.standard.sensors.python import PythonSensor

from airflow_pydantic.tests.conftest import foo

with DAG(
    description="",
    schedule="* * * * *",
    start_date=datetime.fromisoformat("2025-01-01T00:00:00"),
    end_date=datetime.fromisoformat("2026-01-01T00:00:00"),
    max_active_tasks=1,
    max_active_runs=1,
    catchup=False,
    is_paused_upon_creation=True,
    tags=["a", "b"],
    dag_display_name="test",
    dag_id="a-dag",
    default_args={
        "owner": "airflow",
        "email": ["test@test.com"],
        "email_on_failure": True,
        "email_on_retry": True,
        "retries": 3,
        "retry_delay": timedelta(seconds=300.0),
        "start_date": datetime.fromisoformat("2025-01-01T00:00:00"),
        "end_date": datetime.fromisoformat("2026-01-01T00:00:00"),
        "depends_on_past": True,
        "queue": "default",
        "pool": "default",
        "pool_slots": 1,
        "do_xcom_push": True,
        "task_display_name": "test",
    },
) as dag:
    task1 = PythonOperator(
        pool=Pool.get_pool("test-pool-model").pool,
        python_callable=foo,
        op_args=["test"],
        op_kwargs={"test": "test"},
        templates_dict={"test": "test"},
        templates_exts=[".sql", ".hql"],
        show_return_value_in_logs=True,
        task_id="test-python-operator",
        dag=dag,
    )
    task2 = BashOperator(
        pool=Pool.create_or_update_pool(name="test", slots=5, description="", include_deferred=False).pool,
        bash_command="test",
        env={"test": "test"},
        append_env=True,
        output_encoding="utf-8",
        skip_on_exit_code=99,
        cwd="test",
        output_processor=foo,
        task_id="test-bash-operator",
        dag=dag,
    )
    task3 = SSHOperator(
        pool=Pool.get_pool("blerg").pool,
        do_xcom_push=True,
        ssh_hook=SSHHook(remote_host="test", username="test"),
        ssh_conn_id="test",
        command="bash -lc 'set -ex\\ntest1\\ntest2'",
        cmd_timeout=10,
        environment={"test": "test"},
        get_pty=True,
        task_id="test-ssh-operator",
        dag=dag,
    )
    task4 = BashSensor(bash_command="test", task_id="test-bash-sensor", dag=dag)
    task5 = PythonSensor(python_callable=foo, op_args=["test"], op_kwargs={"test": "test"}, task_id="test-python-sensor", dag=dag)
    task1 >> task2
    task1 >> task3
    task2 >> task3
"""
        )
        if _HAVE_AIRFLOW_SSH:
            with pools(Pool()), patch("airflow.models.pool.Pool.get_pool") as mock_get_pool:
                mock_get_pool.return_value = Pool()
                dag.instantiate()
                exec(dag.render())

    def test_render_with_externals(self, dag_with_external):
        assert isinstance(dag_with_external, Dag)
        assert (
            dag_with_external.render(debug_filename="airflow_pydantic/tests/core/rendered/externals.py")
            == """# Generated by airflow-config
from datetime import datetime, timedelta

from airflow.models import DAG
from airflow.models.pool import Pool
from airflow.providers.ssh.operators.ssh import SSHOperator
from airflow.providers.standard.operators.bash import BashOperator
from airflow.providers.standard.operators.python import PythonOperator
from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator

from airflow_pydantic.tests.conftest import foo, hook

with DAG(
    description="",
    schedule="* * * * *",
    start_date=datetime.fromisoformat("2025-01-01T00:00:00"),
    end_date=datetime.fromisoformat("2026-01-01T00:00:00"),
    max_active_tasks=1,
    max_active_runs=1,
    catchup=False,
    is_paused_upon_creation=True,
    tags=["a", "b"],
    dag_display_name="test",
    dag_id="a-dag",
    default_args={
        "owner": "airflow",
        "email": ["test@test.com"],
        "email_on_failure": True,
        "email_on_retry": True,
        "retries": 3,
        "retry_delay": timedelta(seconds=300.0),
        "start_date": datetime.fromisoformat("2025-01-01T00:00:00"),
        "end_date": datetime.fromisoformat("2026-01-01T00:00:00"),
        "depends_on_past": True,
        "queue": "default",
        "pool": "default",
        "pool_slots": 1,
        "do_xcom_push": True,
        "task_display_name": "test",
    },
) as dag:
    task1 = PythonOperator(
        pool=Pool.get_pool("test-pool-model").pool,
        python_callable=foo,
        op_args=["test"],
        op_kwargs={"test": "test"},
        templates_dict={"test": "test"},
        templates_exts=[".sql", ".hql"],
        show_return_value_in_logs=True,
        task_id="test-python-operator",
        dag=dag,
    )
    task2 = BashOperator(
        pool=Pool.create_or_update_pool(name="test", slots=5, description="", include_deferred=False).pool,
        bash_command="test",
        env={"test": "test"},
        append_env=True,
        output_encoding="utf-8",
        skip_on_exit_code=99,
        cwd="test",
        output_processor=foo,
        task_id="test-bash-operator",
        dag=dag,
    )
    task3 = SSHOperator(
        pool=Pool.get_pool("blerg").pool,
        do_xcom_push=True,
        ssh_hook=hook(),
        ssh_conn_id="test",
        command="bash -lc 'set -ex\\ntest1\\ntest2'",
        cmd_timeout=10,
        environment={"test": "test"},
        get_pty=True,
        task_id="test-ssh-operator",
        dag=dag,
    )
    task4 = TriggerDagRunOperator(trigger_dag_id="test_dag", task_id="test-trigger-dagrun", dag=dag)
"""
        )
        if _HAVE_AIRFLOW_SSH:
            with pools(Pool()), patch("airflow.models.pool.Pool.get_pool") as mock_get_pool:
                mock_get_pool.return_value = Pool()
                dag_with_external.instantiate()
                exec(dag_with_external.render())

    def test_render_with_external_supervisor_config(self, dag_with_supervisor):
        assert isinstance(dag_with_supervisor, Dag)
        assert (
            dag_with_supervisor.render(debug_filename="airflow_pydantic/tests/core/rendered/supervisor.py")
            == """# Generated by airflow-config
from datetime import datetime, time, timedelta
from pathlib import Path

from airflow.models import DAG
from airflow_supervisor.airflow.local import Supervisor

with DAG(
    description="",
    schedule="* * * * *",
    start_date=datetime.fromisoformat("2025-01-01T00:00:00"),
    end_date=datetime.fromisoformat("2026-01-01T00:00:00"),
    max_active_tasks=1,
    max_active_runs=1,
    catchup=False,
    is_paused_upon_creation=True,
    tags=["a", "b"],
    dag_display_name="test",
    dag_id="a-dag",
    default_args={
        "owner": "airflow",
        "email": ["test@test.com"],
        "email_on_failure": True,
        "email_on_retry": True,
        "retries": 3,
        "retry_delay": timedelta(seconds=300.0),
        "start_date": datetime.fromisoformat("2025-01-01T00:00:00"),
        "end_date": datetime.fromisoformat("2026-01-01T00:00:00"),
        "depends_on_past": True,
        "queue": "default",
        "pool": "default",
        "pool_slots": 1,
        "do_xcom_push": True,
        "task_display_name": "test",
    },
) as dag:
    task = Supervisor(
        cfg={
            "inet_http_server": {"port": "*:1234", "username": None, "password": None},
            "program": {
                "test": {
                    "command": "bash -c 'sleep 60; exit 1'",
                    "autostart": False,
                    "startsecs": 1,
                    "startretries": None,
                    "autorestart": False,
                    "exitcodes": [0],
                    "stopsignal": "TERM",
                    "stopwaitsecs": 30,
                    "stopasgroup": True,
                    "killasgroup": True,
                    "stdout_logfile": Path("/an/arbitrary/path/test/output.log"),
                    "stderr_logfile": Path("/an/arbitrary/path/test/error.log"),
                    "directory": Path("/an/arbitrary/path/test"),
                }
            },
            "rpcinterface": {"supervisor": {"rpcinterface_factory": "supervisor.rpcinterface:make_main_rpcinterface"}},
            "config_path": Path("/an/arbitrary/path/supervisord.conf"),
            "working_dir": Path("/an/arbitrary/path"),
            "port": "*:1234",
            "endtime": time(23, 59, 0, 0),
        },
        task_id="test-supervisor",
        dag=dag,
    )
"""
        )
        if _HAVE_AIRFLOW:
            dag_with_supervisor.instantiate()
            exec(dag_with_supervisor.render())

    def test_render_with_external_supervisor_ssh_config(self, dag_with_supervisor_ssh):
        from airflow_pydantic.testing import variables

        assert isinstance(dag_with_supervisor_ssh, Dag)
        assert (
            dag_with_supervisor_ssh.render(debug_filename="airflow_pydantic/tests/core/rendered/supervisor_ssh.py")
            == """# Generated by airflow-config
from datetime import datetime, timedelta
from pathlib import Path

from airflow.models import DAG
from airflow_supervisor.airflow.ssh import SupervisorSSH

from airflow_pydantic import Host, Port, Variable

with DAG(
    description="",
    schedule="* * * * *",
    start_date=datetime.fromisoformat("2025-01-01T00:00:00"),
    end_date=datetime.fromisoformat("2026-01-01T00:00:00"),
    max_active_tasks=1,
    max_active_runs=1,
    catchup=False,
    is_paused_upon_creation=True,
    tags=["a", "b"],
    dag_display_name="test",
    dag_id="a-dag",
    default_args={
        "owner": "airflow",
        "email": ["test@test.com"],
        "email_on_failure": True,
        "email_on_retry": True,
        "retries": 3,
        "retry_delay": timedelta(seconds=300.0),
        "start_date": datetime.fromisoformat("2025-01-01T00:00:00"),
        "end_date": datetime.fromisoformat("2026-01-01T00:00:00"),
        "depends_on_past": True,
        "queue": "default",
        "pool": "default",
        "pool_slots": 1,
        "do_xcom_push": True,
        "task_display_name": "test",
    },
) as dag:
    task = SupervisorSSH(
        cfg={
            "inet_http_server": {"port": "*:9001", "username": None, "password": None},
            "program": {
                "test": {
                    "command": "bash -c 'sleep 60; exit 1'",
                    "autostart": False,
                    "startsecs": 1,
                    "startretries": None,
                    "autorestart": False,
                    "exitcodes": [0],
                    "stopsignal": "TERM",
                    "stopwaitsecs": 30,
                    "stopasgroup": True,
                    "killasgroup": True,
                    "stdout_logfile": Path("/an/arbitrary/path/test/output.log"),
                    "stderr_logfile": Path("/an/arbitrary/path/test/error.log"),
                    "directory": Path("/an/arbitrary/path/test"),
                }
            },
            "rpcinterface": {"supervisor": {"rpcinterface_factory": "supervisor.rpcinterface:make_main_rpcinterface"}},
            "config_path": Path("/an/arbitrary/path/supervisord.conf"),
            "working_dir": Path("/an/arbitrary/path"),
        },
        host=Host(name="test_host", username="test_user", password=Variable(key="VAR", deserialize_json=True)),
        port=Port(
            name="test_port", host=Host(name="test_host", username="test_user", password=Variable(key="VAR", deserialize_json=True)), port=8080
        ),
        task_id="test-supervisor",
        dag=dag,
    )
"""
        )
        if _HAVE_AIRFLOW:
            with variables({"user": "test", "password": "password"}):
                dag_with_supervisor_ssh.instantiate()
                exec(dag_with_supervisor_ssh.render())

    def test_render_single_task_python_sensor_and_params(self, dag, python_sensor):
        dag.tasks = {"test_python_sensor": python_sensor}
        dag.params = python_sensor
        assert isinstance(dag, Dag)
        assert isinstance(dag, Dag)
        assert (
            dag.render(debug_filename="airflow_pydantic/tests/core/rendered/sensor.py")
            == """# Generated by airflow-config
from datetime import datetime, timedelta

from airflow.models import DAG
from airflow.models.param import Param
from airflow.providers.standard.sensors.python import PythonSensor

from airflow_pydantic.tests.conftest import foo

with DAG(
    description="",
    schedule="* * * * *",
    start_date=datetime.fromisoformat("2025-01-01T00:00:00"),
    end_date=datetime.fromisoformat("2026-01-01T00:00:00"),
    max_active_tasks=1,
    max_active_runs=1,
    catchup=False,
    params={
        "poke_interval": Param(None, title="Poke Interval", description=None, type=["null", "number"]),
        "timeout": Param(None, title="Timeout", description=None, type=["null", "number"]),
        "soft_fail": Param(None, title="Soft Fail", description=None, type=["null", "boolean"]),
        "mode": Param(None, title="Mode", description=None, type=["null", "string"]),
        "exponential_backoff": Param(None, title="Exponential Backoff", description=None, type=["null", "boolean"]),
        "max_wait": Param(None, title="Max Wait", description=None, type=["null", "number"]),
        "silent_fail": Param(None, title="Silent Fail", description=None, type=["null", "boolean"]),
        "never_fail": Param(None, title="Never Fail", description=None, type=["null", "boolean"]),
        "op_args": Param(
            ["test"], title="Op Args", description="a list of positional arguments that will get unpacked when calling your callable", type="array"
        ),
        "op_kwargs": Param(
            {"test": "test"},
            title="Op Kwargs",
            description="a dictionary of keyword arguments that will get unpacked in your function",
            type="object",
        ),
        "templates_dict": Param(
            None,
            title="Templates Dict",
            description="a dictionary where the values are templates that will get templated by the Airflow engine sometime between __init__ and execute takes place and are made available in your callable’s context after the template has been applied. (templated)",
            type=["null", "object"],
        ),
    },
    is_paused_upon_creation=True,
    tags=["a", "b"],
    dag_display_name="test",
    dag_id="a-dag",
    default_args={
        "owner": "airflow",
        "email": ["test@test.com"],
        "email_on_failure": True,
        "email_on_retry": True,
        "retries": 3,
        "retry_delay": timedelta(seconds=300.0),
        "start_date": datetime.fromisoformat("2025-01-01T00:00:00"),
        "end_date": datetime.fromisoformat("2026-01-01T00:00:00"),
        "depends_on_past": True,
        "queue": "default",
        "pool": "default",
        "pool_slots": 1,
        "do_xcom_push": True,
        "task_display_name": "test",
    },
) as dag:
    test_python_sensor = PythonSensor(python_callable=foo, op_args=["test"], op_kwargs={"test": "test"}, task_id="test-python-sensor", dag=dag)
"""
        )
        # Can't eval due to null param
        # if _HAVE_AIRFLOW:
        #     exec(dag.render())

    def test_render_dag_with_attribute_dependencies(self, dag_with_attribute_dependencies):
        assert (
            dag_with_attribute_dependencies.render(debug_filename="airflow_pydantic/tests/core/rendered/attr.py")
            == """# Generated by airflow-config
from datetime import datetime, timedelta

from airflow.models import DAG
from airflow.providers.standard.operators.python import PythonOperator
from airflow_ha.operator import HighAvailabilitySensor

from airflow_pydantic.tests.conftest import _choose, _null, fail

with DAG(
    description="",
    schedule="* * * * *",
    start_date=datetime.fromisoformat("2025-01-01T00:00:00"),
    end_date=datetime.fromisoformat("2026-01-01T00:00:00"),
    max_active_tasks=1,
    max_active_runs=1,
    catchup=False,
    is_paused_upon_creation=True,
    tags=["a", "b"],
    dag_display_name="test",
    dag_id="a-dag",
    default_args={
        "owner": "airflow",
        "email": ["test@test.com"],
        "email_on_failure": True,
        "email_on_retry": True,
        "retries": 3,
        "retry_delay": timedelta(seconds=300.0),
        "start_date": datetime.fromisoformat("2025-01-01T00:00:00"),
        "end_date": datetime.fromisoformat("2026-01-01T00:00:00"),
        "depends_on_past": True,
        "queue": "default",
        "pool": "default",
        "pool_slots": 1,
        "do_xcom_push": True,
        "task_display_name": "test",
    },
) as dag:
    pre = PythonOperator(python_callable=_null, task_id="pre", dag=dag)
    ha = HighAvailabilitySensor(python_callable=_choose, task_id="ha", dag=dag)
    retrigger_fail = PythonOperator(python_callable=_null, task_id="retrigger_fail", dag=dag)
    stop_fail = PythonOperator(trigger_rule="all_failed", python_callable=fail, task_id="stop_fail", dag=dag)
    retrigger_pass = PythonOperator(python_callable=_null, task_id="retrigger_pass", dag=dag)
    stop_pass = PythonOperator(python_callable=_null, task_id="lam-stop_pass", dag=dag)
    pre >> ha
    ha.retrigger_fail >> retrigger_fail
    ha.stop_fail >> stop_fail
    ha.retrigger_pass >> retrigger_pass
    ha.stop_pass >> stop_pass
"""
        )
        if _HAVE_AIRFLOW:
            exec(dag_with_attribute_dependencies.render())
