"""Auto-generated Airflow DAG for {{ contract_id }}.

Generated by Griot scaffold on {{ generated_at }}.
"""
from __future__ import annotations

from datetime import datetime, timedelta

from airflow import DAG
from airflow.operators.python import PythonOperator

from griot_core.integrations.airflow import GriotValidateOperator, GriotReportOperator


default_args = {
    "owner": "{{ owner | default('data-engineering') }}",
    "depends_on_past": False,
    "retries": 1,
    "retry_delay": timedelta(minutes=5),
}


with DAG(
    dag_id="{{ dag_id }}",
    default_args=default_args,
    description="Pipeline for {{ contract_id }} contract",
    schedule_interval="{{ schedule | default('@daily') }}",
    start_date=datetime(2024, 1, 1),
    catchup=False,
    tags=["griot", "{{ contract_id }}"],
) as dag:

    # --Extract ------------------------------------------------------
{%- if source and source != 'none' %}

    def _extract(**kwargs):
        """Extract data from {{ source }} source.

        TODO: Implement extraction logic.
        """
        raise NotImplementedError("Implement {{ source }} extraction")

    extract = PythonOperator(
        task_id="extract",
        python_callable=_extract,
    )
{%- else %}

    def _extract(**kwargs):
        """Extract data from source.

        TODO: Implement extraction logic.
        """
        raise NotImplementedError("Implement extraction")

    extract = PythonOperator(
        task_id="extract",
        python_callable=_extract,
    )
{%- endif %}

    # --Validate with Griot ------------------------------------------

    validate = GriotValidateOperator(
        task_id="validate_contract",
        contract_id="{{ contract_id }}",
        profile="{{ profile | default('default') }}",
        fail_on_invalid=True,
    )

    # --Report results to Griot Registry -----------------------------
    # Using Airflow Connection for registry auth
    # Create a connection in Airflow with:
    #   conn_id: griot_registry
    #   conn_type: http
    #   host: <registry_url>
    #   login: <client_id>
    #   password: <client_secret>
    report = GriotReportOperator(
        task_id="report_to_registry",
        contract_id="{{ contract_id }}",
        griot_conn_id="griot_registry",
        environment="production",
    )

    # --Transform ----------------------------------------------------

    def _transform(**kwargs):
        """Transform data according to contract schema.

        TODO: Implement transformation logic.
        """
        raise NotImplementedError("Implement transformation")

    transform = PythonOperator(
        task_id="transform",
        python_callable=_transform,
    )

    # --Load ---------------------------------------------------------

    def _load(**kwargs):
        """Load data to {{ target | default('target') }}.

        TODO: Implement load logic.
        """
        raise NotImplementedError("Implement load to {{ target | default('target') }}")

    load = PythonOperator(
        task_id="load",
        python_callable=_load,
    )

    # --DAG flow -----------------------------------------------------

    extract >> validate >> report >> transform >> load
