from __future__ import annotations

from airflow.sdk import DAG, task

from orchestration.airflow.engine.local_dev_assets import (
    load_hello_asset_smoke_config,
)


HELLO_ASSETS = load_hello_asset_smoke_config()
DEFAULT_ARGS = {"owner": HELLO_ASSETS.owner}


with DAG(
    dag_id=HELLO_ASSETS.producer_dag_id,
    default_args=DEFAULT_ARGS,
    schedule=HELLO_ASSETS.producer_schedule,
    start_date=HELLO_ASSETS.start_date,
    catchup=False,
    tags=HELLO_ASSETS.tags,
) as producer_dag:

    @task(outlets=[HELLO_ASSETS.raw_asset])
    def emit_hello_asset() -> dict[str, object]:
        return {"asset": HELLO_ASSETS.raw_asset_uri, "row_count": 1}

    emit_hello_asset()


with DAG(
    dag_id=HELLO_ASSETS.consumer_dag_id,
    default_args=DEFAULT_ARGS,
    schedule=HELLO_ASSETS.consumer_schedule,
    start_date=HELLO_ASSETS.start_date,
    catchup=False,
    tags=HELLO_ASSETS.tags,
) as consumer_dag:

    @task(inlets=[HELLO_ASSETS.raw_asset], outlets=[HELLO_ASSETS.consumed_asset])
    def consume_hello_asset() -> dict[str, object]:
        return {
            "input_asset": HELLO_ASSETS.raw_asset_uri,
            "output_asset": HELLO_ASSETS.consumed_asset_uri,
            "row_count": 1,
        }

    consume_hello_asset()