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

import pytest
from airflow_balancer.testing import pools
from hydra.errors import InstantiationException

from airflow_config import load_config

RENDERED_DAG = """# Generated by airflow-config
from airflow.models import DAG
from airflow.operators.bash import BashOperator
from airflow.operators.python import PythonOperator
from airflow.operators.python import ShortCircuitOperator
from airflow.providers.ssh.hooks.ssh import SSHHook
from airflow.providers.ssh.operators.ssh import SSHOperator
from airflow_config.tests.conftest import should_short_circuit
from airflow_config.tests.setups.utils.print_hello import print_hello
with DAG('example_dag') as dag:
    task_1 = SSHOperator(do_xcom_push=True, ssh_hook=SSHHook(remote_host='test', username='test', port=22, cmd_timeout=10, keepalive_interval=30, banner_timeout=30.0), 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'", append_env=False, output_encoding='utf-8', skip_on_exit_code=99, task_id='task_3', dag=dag)
    task_4 = BashOperator(bash_command='echo `pwd`', append_env=False, output_encoding='utf-8', skip_on_exit_code=99, task_id='task_4', dag=dag)
    task_5 = PythonOperator(python_callable=print_hello, op_args=[], op_kwargs={}, templates_dict={}, 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 airflow.models import DAG
from airflow.providers.ssh.hooks.ssh import SSHHook
from airflow.providers.ssh.operators.ssh import SSHOperator
with DAG('example_dag') as dag:
    task_1 = SSHOperator(ssh_hook=SSHHook(remote_host='test', username='test', port=22, cmd_timeout=10, keepalive_interval=30, banner_timeout=30.0), ssh_conn_id='test', command='test', task_id='task_1', dag=dag)
"""


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"]
        (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"]
        (Path(tmp_dir) / "example_dag.py").read_text() == RENDERED_DAG_MULTI


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():
        cfg = load_config("config", "balancer")
        assert cfg.dags["example_dag"].tasks["task_1"].ssh_hook.remote_host == "server2.local"


def test_render_self_reference():
    # TODO: fix this
    with pytest.raises(InstantiationException):
        load_config("config", "self_reference")
