from os import listdir
from pathlib import Path
from tempfile import TemporaryDirectory

from airflow_balancer.testing import pools, variables

from airflow_config import load_config

RENDERED_DAG = """# Generated by airflow-config
from datetime import datetime

from airflow.models import DAG
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, ShortCircuitOperator

from airflow_config.tests.conftest import should_short_circuit
from airflow_config.tests.setups.utils.print_hello import print_hello

with DAG(
    description="this is an example dag",
    schedule="0 3 * * *",
    start_date=datetime.fromisoformat("2024-01-01T00:00:00"),
    catchup=False,
    tags=["utility", "test"],
    dag_id="example_dag",
    default_args={
        "owner": "custom_owner",
        "email": ["myemail@myemail.com"],
        "email_on_failure": False,
        "email_on_retry": False,
        "retries": 0,
        "depends_on_past": False,
    },
) as dag:
    task_1 = SSHOperator(
        do_xcom_push=True,
        ssh_hook=SSHHook(remote_host="test", username="test"),
        ssh_conn_id="test",
        command="test",
        cmd_timeout=10,
        environment={"test": "test"},
        get_pty=True,
        task_id="task_1",
        dag=dag,
    )
    task_2 = ShortCircuitOperator(python_callable=should_short_circuit, task_id="task_2", dag=dag)
    task_3 = BashOperator(bash_command="echo '1'", task_id="task_3", dag=dag)
    task_4 = BashOperator(bash_command="echo `pwd`", task_id="task_4", dag=dag)
    task_5 = PythonOperator(
        python_callable=print_hello,
        op_args=[],
        op_kwargs={},
        templates_dict={},
        templates_exts=None,
        show_return_value_in_logs=True,
        task_id="task_5",
        dag=dag,
    )
    task_1 >> task_2
    task_2 >> task_3
    task_3 >> task_4
    task_4 >> task_5
"""

RENDERED_DAG_MULTI = """# Generated by airflow-config
from datetime import datetime

from airflow.models import DAG
from airflow.providers.ssh.hooks.ssh import SSHHook
from airflow.providers.ssh.operators.ssh import SSHOperator

with DAG(
    schedule="0 3 * * *",
    start_date=datetime.fromisoformat("2024-01-01T00:00:00"),
    catchup=False,
    tags=["utility", "test"],
    dag_id="example_dag",
    default_args={},
) as dag:
    task_1 = SSHOperator(ssh_hook=SSHHook(remote_host="test", username="test"), ssh_conn_id="test", command="test", task_id="task_1", dag=dag)
"""

RENDERED_DAG_MULTI2 = """# Generated by airflow-config
from datetime import datetime

from airflow.models import DAG
from airflow.providers.standard.operators.bash import BashOperator

with DAG(
    schedule="0 4 * * *",
    start_date=datetime.fromisoformat("2024-01-01T00:00:00"),
    catchup=False,
    tags=["utility", "test"],
    dag_id="example_dag2",
    default_args={},
) as dag:
    task_1 = BashOperator(bash_command='echo "1"', task_id="task_1", dag=dag)
"""

RENDERED_DAG_TYPE_OVERRIDE = """# Generated by airflow-config
from datetime import datetime

from airflow.models import DAG
from airflow.providers.standard.operators.bash import BashOperator
from airflow.providers.standard.operators.python import PythonOperator

from airflow_config.tests.setups.utils.print_hello import print_hello

with DAG(
    description="this is an example dag",
    schedule="0 3 * * *",
    start_date=datetime.fromisoformat("2024-01-01T00:00:00"),
    catchup=False,
    tags=["utility", "test"],
    dag_id="example_dag",
    default_args={
        "owner": "custom_owner",
        "email": ["myemail@myemail.com"],
        "email_on_failure": False,
        "email_on_retry": False,
        "retries": 0,
        "depends_on_past": False,
    },
) as dag:
    task_1 = BashOperator(bash_command="echo `pwd`", task_id="task_1", dag=dag)
    task_2 = PythonOperator(python_callable=print_hello, task_id="task_2", dag=dag)
"""

RENDERED_DAG_TYPE_OVERRIDE2 = """# Generated by airflow-config
from datetime import datetime

from airflow.models import DAG
from airflow.providers.standard.operators.bash import BashOperator

with DAG(
    schedule="0 4 * * *",
    start_date=datetime.fromisoformat("2024-01-01T00:00:00"),
    catchup=False,
    tags=["utility", "test"],
    dag_id="example_dag2",
    default_args={
        "owner": "custom_owner2",
        "email": ["myemail@myemail.com"],
        "email_on_failure": False,
        "email_on_retry": False,
        "retries": 0,
        "depends_on_past": False,
    },
) as dag:
    task_1 = BashOperator(bash_command='echo "1"', task_id="task_1", dag=dag)
    task_2 = BashOperator(bash_command="echo 2", task_id="task_2", dag=dag)
    task_3 = BashOperator(bash_command="echo 3", task_id="task_3", dag=dag)
"""

RENDERED_DAG_BALANCER = """# Generated by airflow-config
from airflow.models import DAG
from airflow.models.pool import Pool
from airflow.models.variable import Variable as AirflowVariable
from airflow.providers.ssh.hooks.ssh import SSHHook
from airflow.providers.ssh.operators.ssh import SSHOperator

with DAG(schedule="0 3 * * *", dag_id="example_dag", default_args={}) as dag:
    task_1 = SSHOperator(
        pool=Pool.create_or_update_pool(name="server2", slots=8, description="Balancer pool for host(server2)", include_deferred=False).pool,
        ssh_hook=SSHHook(remote_host="server2.local", username="user1", password=AirflowVariable.get("myvar", deserialize_json=True)["password"]),
        ssh_conn_id="test",
        command="test",
        task_id="task_1",
        dag=dag,
    )
"""

RENDERED_DAG_SELF_REFERENCE = """# Generated by airflow-config
from airflow.models import DAG
from airflow.providers.ssh.operators.ssh import SSHOperator

from airflow_config.tests.setups.factory.utils import get_ssh_hook_from_config

with DAG(schedule="0 3 * * *", dag_id="example_dag", default_args={}) as dag:
    task_1 = SSHOperator(ssh_hook=get_ssh_hook_from_config(), ssh_conn_id="test", command="test", task_id="task_1", dag=dag)
"""

RENDERED_DAG_TEMPLATES = """# Generated by airflow-config
from datetime import datetime

from airflow.models import DAG
from airflow.providers.ssh.hooks.ssh import SSHHook
from airflow.providers.ssh.operators.ssh import SSHOperator
from airflow.providers.standard.operators.python import PythonOperator

from airflow_config.tests.setups.utils.print_hello import print_hello

with DAG(
    description="this is an example dag",
    schedule="0 3 * * *",
    start_date=datetime.fromisoformat("2024-01-01T00:00:00"),
    catchup=False,
    tags=["utility", "test"],
    dag_id="example_dag",
    default_args={
        "owner": "custom_owner",
        "email": ["myemail@myemail.com"],
        "email_on_failure": False,
        "email_on_retry": False,
        "retries": 0,
        "depends_on_past": False,
    },
) as dag:
    task_1 = PythonOperator(
        python_callable=print_hello,
        op_args=[],
        op_kwargs={},
        templates_dict={},
        templates_exts=None,
        show_return_value_in_logs=True,
        task_id="task_1",
        dag=dag,
    )
    task_2 = SSHOperator(
        do_xcom_push=True,
        ssh_hook=SSHHook(remote_host="test", username="test"),
        ssh_conn_id="test",
        command="test",
        cmd_timeout=10,
        environment={"blerg": "blerg", "test": "test"},
        get_pty=True,
        task_id="task_2",
        dag=dag,
    )
    task_1 >> task_2
"""


def test_render():
    conf = load_config("config", "factory")
    assert conf.dags["example_dag"].render() == RENDERED_DAG


def test_generate():
    conf = load_config("config", "factory")
    with TemporaryDirectory() as tmp_dir:
        conf.generate(tmp_dir)
        assert sorted(listdir(tmp_dir)) == ["example_dag.py", "example_dag2.py"]
        assert (Path(tmp_dir) / "example_dag.py").read_text() == RENDERED_DAG


def test_render_multi():
    conf = load_config("config", "multi")
    assert conf.dags["example_dag"].render() == RENDERED_DAG_MULTI


def test_generate_multi():
    conf = load_config("config", "multi")
    with TemporaryDirectory() as tmp_dir:
        conf.generate(tmp_dir)
        assert sorted(listdir(tmp_dir)) == ["example_dag.py", "example_dag2.py"]
        assert (Path(tmp_dir) / "example_dag.py").read_text() == RENDERED_DAG_MULTI
        assert (Path(tmp_dir) / "example_dag2.py").read_text() == RENDERED_DAG_MULTI2


def test_generate_type_override():
    conf = load_config("config", "type_override")
    with TemporaryDirectory() as tmp_dir:
        conf.generate(tmp_dir)
        assert sorted(listdir(tmp_dir)) == ["example_dag.py", "example_dag2.py"]
        assert (Path(tmp_dir) / "example_dag.py").read_text() == RENDERED_DAG_TYPE_OVERRIDE
        assert (Path(tmp_dir) / "example_dag2.py").read_text() == RENDERED_DAG_TYPE_OVERRIDE2


def test_render_dynamic_gen(has_airflow):
    conf = load_config("config", "factory")
    assert conf.default_dag_args.start_date is not None
    conf.generate_in_mem()


def test_render_self_host():
    cfg = load_config("config", "host")
    assert cfg.dags["example_dag"].tasks["task_1"].ssh_hook.remote_host == "test_host.local"


def test_render_self_balancer_query():
    with pools(), variables({"password": "thepassword"}):
        cfg = load_config("config", "balancer")
        assert cfg.dags["example_dag"].tasks["task_1"].ssh_hook.remote_host == "server2.local"
        with TemporaryDirectory() as tmp_dir:
            cfg.generate(tmp_dir)
            assert sorted(listdir(tmp_dir)) == ["example_dag.py"]
            assert (Path(tmp_dir) / "example_dag.py").read_text() == RENDERED_DAG_BALANCER


def test_render_self_reference():
    cfg = load_config("config", "self_reference")
    with TemporaryDirectory() as tmp_dir:
        cfg.generate(tmp_dir)
        assert sorted(listdir(tmp_dir)) == ["example_dag.py"]
        assert (Path(tmp_dir) / "example_dag.py").read_text() == RENDERED_DAG_SELF_REFERENCE


def test_render_templates():
    cfg = load_config("config", "templates")
    with TemporaryDirectory() as tmp_dir:
        cfg.generate(tmp_dir)
        assert sorted(listdir(tmp_dir)) == ["example_dag.py"]
        assert (Path(tmp_dir) / "example_dag.py").read_text() == RENDERED_DAG_TEMPLATES
