from datetime import datetime, timedelta
from os import listdir
from pathlib import Path
from tempfile import TemporaryDirectory

from airflow.timetables.interval import DeltaDataIntervalTimetable

from airflow_config import DAG, create_dag, load_config


def test_config_and_options():
    conf = load_config("config", "test")
    assert conf.default_args.owner == "test"
    assert conf.default_args.email == ["myemail@myemail.com"]
    assert conf.default_args.email_on_failure is False
    assert conf.default_args.email_on_retry is False
    assert conf.default_args.retries == 0
    assert conf.default_args.depends_on_past is False
    # assert conf.global_.schedule == timedelta(seconds=60)
    assert conf.default_dag_args.start_date == datetime(2024, 1, 1)
    assert conf.default_dag_args.catchup is False
    assert conf.default_dag_args.tags == ["utility", "test"]

    assert conf.dags["example_dag"].default_args.owner == "custom_owner"


def test_create_dag_from_config():
    conf = load_config("config", "test")
    d = DAG(dag_id="testdag", config=conf)
    assert d.default_args["owner"] == "test"
    assert d.default_args["email"] == ["myemail@myemail.com"]
    assert d.default_args["email_on_failure"] is False
    assert d.default_args["email_on_retry"] is False
    assert d.default_args["retries"] == 0
    assert d.default_args["depends_on_past"] is False
    assert d.schedule_interval == timedelta(seconds=3600)
    assert isinstance(d.timetable, DeltaDataIntervalTimetable)
    assert isinstance(d.timetable._delta, timedelta)
    assert d.start_date.year == 2024
    assert d.start_date.month == 1
    assert d.start_date.day == 1
    assert d.catchup is False
    assert d.tags == ["utility", "test"]

    d = DAG(dag_id="example_dag", config=conf)
    assert d.default_args["owner"] == "custom_owner"
    assert d.default_args["email"] == ["myemail@myemail.com"]
    assert d.schedule_interval == "0 3 * * *"


def test_create_dag_tasks_from_config():
    conf = load_config("config", "test")
    d = DAG(dag_id="example_dag", config=conf)
    assert len(d.tasks) == 5
    assert d.tasks[0].task_id == "task_1"
    assert d.tasks[1].task_id == "task_2"
    assert d.tasks[1].upstream_task_ids == {"task_1"}
    assert d.tasks[2].task_id == "task_3"
    assert d.tasks[2].upstream_task_ids == {"task_2"}
    assert d.tasks[3].task_id == "task_4"
    assert d.tasks[3].upstream_task_ids == {"task_3"}
    assert d.tasks[4].task_id == "task_5"
    assert d.tasks[4].upstream_task_ids == {"task_4"}
    assert conf.dags["example_dag"].tasks["task_5"].model_dump()["show_return_value_in_logs"] is True


def test_create_dag_from_config_create_dag():
    d = create_dag("config", "test")
    assert d.dag_id == "tests-setups-good-dag-factory-test-dag-factory"
    assert d.dag_id in globals()
    assert d.default_args["owner"] == "test"


def test_serialize_with_airflow_extras():
    conf = load_config("config", "test")
    print(conf.model_dump_json(serialize_as_any=True))
    assert '"operator":"airflow.operators.bash.BashOperator"' in conf.model_dump_json()


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.good.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
"""


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


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