from __future__ import annotations

import pendulum
from airflow.sdk import DAG, task

from orchestration.airflow.engine.extract_assets import (
    REPO_ROOT,
    build_airflow_asset,
    execute_extract_command,
    load_extract_asset_config,
    manifest_event_extra,
)
from orchestration.airflow.engine.tidb_extract_operator import TiDBCloudExportManifestOperator


EXTRACT_ASSETS_CONFIG_PATH = REPO_ROOT / "orchestration" / "airflow" / "sources" / "extract-assets.yaml"
EXTRACT_ASSETS = load_extract_asset_config(EXTRACT_ASSETS_CONFIG_PATH)


def _build_bigquery_extract_dag(unit):
    manifest_asset = build_airflow_asset(unit.manifest_asset)
    source_outlets = [build_airflow_asset(asset) for asset in unit.source_outlets]

    with DAG(
        dag_id=unit.dag_id,
        default_args={"owner": unit.owner},
        schedule=unit.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.generate_contracts_task_id)
        def generate_contracts() -> dict[str, object]:
            return execute_extract_command(unit.command, "generate-contracts")

        @task(task_id=unit.export_task_id, outlets=source_outlets)
        def export_snapshot(_: dict[str, object]) -> dict[str, object]:
            return execute_extract_command(unit.command, "export")

        @task(task_id=unit.write_manifest_task_id, outlets=[manifest_asset])
        def write_manifest(export_result: dict[str, object], *, outlet_events=None) -> dict[str, object]:
            manifest = execute_extract_command(
                unit.command,
                "write-manifest",
                run_date=str(export_result["run_date"]),
                run_id=str(export_result["run_id"]),
            )
            extra = manifest_event_extra(
                manifest_payload=manifest,
                fallback_payload=export_result,
            )
            outlet_events[manifest_asset].extra = extra
            return {"manifest": extra}

        write_manifest(export_snapshot(generate_contracts()))

    return dag


def _build_tidbcloud_extract_dag(unit):
    manifest_asset = build_airflow_asset(unit.manifest_asset)
    source_outlets = [build_airflow_asset(asset) for asset in unit.source_outlets]

    with DAG(
        dag_id=unit.dag_id,
        default_args={"owner": unit.owner},
        schedule=unit.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.generate_contracts_task_id)
        def generate_contracts() -> dict[str, object]:
            return execute_extract_command(unit.command, "generate-contracts")

        manifest = TiDBCloudExportManifestOperator(
            task_id=unit.write_manifest_task_id,
            command=unit.command,
            manifest_asset=manifest_asset,
            poll_interval_seconds=unit.poll_interval_seconds,
            timeout_seconds=unit.timeout_seconds,
            outlets=[*source_outlets, manifest_asset],
        )
        generate_contracts() >> manifest

    return dag


for extract_asset_unit in EXTRACT_ASSETS.units:
    if extract_asset_unit.mode == "bigquery_sync":
        globals()[extract_asset_unit.dag_id] = _build_bigquery_extract_dag(extract_asset_unit)
    elif extract_asset_unit.mode == "tidbcloud_deferrable":
        globals()[extract_asset_unit.dag_id] = _build_tidbcloud_extract_dag(extract_asset_unit)
    else:
        raise ValueError(f"unsupported extract asset mode: {extract_asset_unit.mode!r}")