from __future__ import annotations

import pendulum
from airflow.sdk import DAG, get_current_context, task

from orchestration.airflow.engine.load_assets import (
    REPO_ROOT,
    build_airflow_asset,
    execute_load_asset_command,
    load_load_asset_config,
    resolve_manifest_path,
)


LOAD_ASSETS_CONFIG_PATH = REPO_ROOT / "orchestration" / "airflow" / "sources" / "load-assets.yaml"
LOAD_ASSETS = load_load_asset_config(LOAD_ASSETS_CONFIG_PATH)


def _build_load_asset_dag(unit):
    schedule = (
        [build_airflow_asset(asset) for asset in unit.schedule]
        if isinstance(unit.schedule, tuple)
        else unit.schedule
    )
    inlets = [build_airflow_asset(asset) for asset in unit.inlets]
    outlets = [build_airflow_asset(asset) for asset in unit.outlets]

    with DAG(
        dag_id=unit.dag_id,
        default_args={"owner": unit.owner},
        schedule=schedule,
        start_date=pendulum.parse(unit.start_date),
        catchup=unit.catchup,
        max_active_runs=unit.max_active_runs,
        tags=list(unit.tags),
    ) as dag:

        @task(
            task_id=unit.task_id,
            inlets=inlets,
            outlets=outlets,
        )
        def run_load_asset() -> dict[str, object]:
            manifest_path = None
            if unit.mode == "manifest":
                manifest_path = resolve_manifest_path(unit, context=get_current_context())
            return execute_load_asset_command(
                unit.command,
                output_assets=unit.outlets,
                manifest_path=manifest_path,
            )

        run_load_asset()

    return dag


for load_asset_unit in LOAD_ASSETS.units:
    globals()[load_asset_unit.dag_id] = _build_load_asset_dag(load_asset_unit)