from __future__ import annotations

import pendulum
from pathlib import Path
from airflow.sdk import DAG
from cosmos import DbtTaskGroup
from cosmos.config import ExecutionConfig, ProfileConfig, ProjectConfig, RenderConfig
from cosmos.constants import ExecutionMode, InvocationMode, LoadMode, TestBehavior

from orchestration.airflow.engine.dbt_assets import (
    build_airflow_asset,
    build_contract_asset_converters,
    load_dbt_asset_config,
)


DBT_REVREC_CONFIG_PATH = Path(__file__).resolve().parents[1] / "sources" / "dbt-revrec.yaml"
DBT_REVREC = load_dbt_asset_config(DBT_REVREC_CONFIG_PATH)


with DAG(
    dag_id=DBT_REVREC.dag_id,
    default_args={"owner": DBT_REVREC.owner},
    schedule=[build_airflow_asset(asset) for asset in DBT_REVREC.schedule_assets],
    start_date=pendulum.parse(DBT_REVREC.start_date),
    catchup=DBT_REVREC.catchup,
    max_active_runs=DBT_REVREC.max_active_runs,
    tags=list(DBT_REVREC.tags),
) as dag:
    DbtTaskGroup(
        group_id=DBT_REVREC.task_group_id,
        project_config=ProjectConfig(
            dbt_project_path=DBT_REVREC.project_path,
            models_relative_path=".",
            manifest_path=DBT_REVREC.manifest_path,
            install_dbt_deps=False,
        ),
        profile_config=ProfileConfig(
            profile_name=DBT_REVREC.profile_name,
            target_name=DBT_REVREC.target_name,
            profiles_yml_filepath=DBT_REVREC.profiles_yml_path,
        ),
        execution_config=ExecutionConfig(
            execution_mode=ExecutionMode.LOCAL,
            invocation_mode=InvocationMode.SUBPROCESS,
            dbt_executable_path=DBT_REVREC.dbt_executable_path,
        ),
        render_config=RenderConfig(
            emit_datasets=False,
            load_method=LoadMode.DBT_MANIFEST,
            invocation_mode=InvocationMode.SUBPROCESS,
            select=list(DBT_REVREC.select),
            test_behavior=TestBehavior.AFTER_EACH,
            node_converters=build_contract_asset_converters(DBT_REVREC.node_bindings),
            node_conversion_by_task_group=False,
        ),
        operator_args={
            "append_env": True,
        },
    )