import pytest

from airflow_pydantic.migration import _airflow_3

dags = {
    "test_one": """# Generated by airflow-config
from airflow.models import DAG
from airflow.providers.standard.operators.bash import BashOperator
from pendulum import Timezone, datetime as pdatetime

with DAG(
    schedule="0 0 * * *",
    start_date=pdatetime(year=2025, month=1, day=1, tz=Timezone("America/New_York")),
    max_active_runs=1,
    catchup=False,
    dag_id="test_one",
    default_args={},
) as dag:
    one = BashOperator(bash_command="echo 'one'", task_id="one", dag=dag)
""",
    "test_two": """# Generated by airflow-config
from airflow.models import DAG
from airflow.providers.standard.operators.bash import BashOperator
from airflow.timetables.interval import CronDataIntervalTimetable
from pendulum import Timezone, datetime as pdatetime

with DAG(
    schedule=CronDataIntervalTimetable(cron="0 0 * * *", timezone="America/New_York"),
    start_date=pdatetime(year=2025, month=1, day=1, tz=Timezone("America/New_York")),
    max_active_runs=1,
    catchup=False,
    dag_id="test_two",
    default_args={},
) as dag:
    two = BashOperator(bash_command="echo 'two'", task_id="two", dag=dag)
""",
    "test_three": """# Generated by airflow-config
from airflow.models import DAG
from airflow.providers.standard.operators.bash import BashOperator
from airflow.timetables.trigger import CronTriggerTimetable
from dateutil.relativedelta import relativedelta
from pendulum import Timezone, datetime as pdatetime

with DAG(
    schedule=CronTriggerTimetable(
        cron="0 0 * * *",
        timezone="America/New_York",
        interval=relativedelta(years=0, months=0, days=0, leapdays=0, hours=5, minutes=10, seconds=0, microseconds=0),
    ),
    start_date=pdatetime(year=2025, month=1, day=1, tz=Timezone("America/New_York")),
    max_active_runs=1,
    catchup=False,
    dag_id="test_three",
    default_args={},
) as dag:
    three = BashOperator(bash_command="echo 'three'", task_id="three", dag=dag)
""",
    "test_four": """# Generated by airflow-config
from airflow.models import DAG
from airflow.providers.standard.operators.bash import BashOperator
from airflow.timetables.trigger import MultipleCronTriggerTimetable
from dateutil.relativedelta import relativedelta
from pendulum import Timezone, datetime as pdatetime

with DAG(
    schedule=MultipleCronTriggerTimetable(
        "0 0 * * *",
        "0 1 * * *",
        timezone="America/New_York",
        interval=relativedelta(years=0, months=0, days=0, leapdays=0, hours=5, minutes=10, seconds=0, microseconds=0),
    ),
    start_date=pdatetime(year=2025, month=1, day=1, tz=Timezone("America/New_York")),
    max_active_runs=1,
    catchup=False,
    dag_id="test_four",
    default_args={},
) as dag:
    four = BashOperator(bash_command="echo 'four'", task_id="four", dag=dag)
""",
    "test_five": """# Generated by airflow-config
from airflow.models import DAG
from airflow.providers.standard.operators.bash import BashOperator
from airflow.timetables.interval import DeltaDataIntervalTimetable
from dateutil.relativedelta import relativedelta
from pendulum import Timezone, datetime as pdatetime

with DAG(
    schedule=DeltaDataIntervalTimetable(delta=relativedelta(years=0, months=0, days=1, leapdays=0, hours=2, minutes=3, seconds=0, microseconds=0)),
    start_date=pdatetime(year=2025, month=1, day=1, tz=Timezone("America/New_York")),
    max_active_runs=1,
    catchup=False,
    dag_id="test_five",
    default_args={},
) as dag:
    five = BashOperator(bash_command="echo 'five'", task_id="five", dag=dag)
""",
    "test_six": """# Generated by airflow-config
from airflow.models import DAG
from airflow.providers.standard.operators.bash import BashOperator
from airflow.timetables.trigger import DeltaTriggerTimetable
from dateutil.relativedelta import relativedelta
from pendulum import Timezone, datetime as pdatetime

with DAG(
    schedule=DeltaTriggerTimetable(
        delta=relativedelta(years=0, months=0, days=1, leapdays=0, hours=2, minutes=3, seconds=0, microseconds=0),
        interval=relativedelta(years=0, months=0, days=0, leapdays=0, hours=5, minutes=10, seconds=0, microseconds=0),
    ),
    start_date=pdatetime(year=2025, month=1, day=1, tz=Timezone("America/New_York")),
    max_active_runs=1,
    catchup=False,
    dag_id="test_six",
    default_args={},
) as dag:
    six = BashOperator(bash_command="echo 'six'", task_id="six", dag=dag)
""",
    "test_seven": """# Generated by airflow-config
from datetime import datetime

from airflow.models import DAG
from airflow.providers.standard.operators.bash import BashOperator
from airflow.timetables.events import EventsTimetable
from pendulum import Timezone, datetime as pdatetime

with DAG(
    schedule=EventsTimetable(
        event_dates=[datetime.fromisoformat("2025-01-01T00:00:00"), pdatetime(year=2025, month=1, day=2, tz=Timezone("America/New_York"))],
        restrict_to_events=True,
        presorted=True,
        description="whatever",
    ),
    start_date=pdatetime(year=2025, month=1, day=1, tz=Timezone("America/New_York")),
    max_active_runs=1,
    catchup=False,
    dag_id="test_seven",
    default_args={},
) as dag:
    seven = BashOperator(bash_command="echo 'seven'", task_id="seven", dag=dag)
""",
}


class TestConfig:
    @pytest.mark.parametrize("dag_id", list(dags.keys()))
    def test_dag_render(self, load_config, dag_id):
        conf = load_config("config", "config")

        assert conf.dags[dag_id].render() == dags[dag_id]

        if _airflow_3() is None:
            # Skip
            return

        if _airflow_3() is False and dag_id in ("test_four", "test_six"):
            # Skip, not available on airflow 2
            return

        exec(conf.dags[dag_id].render())
