from airflow import DAG
from airflow.operators.bash_operator import BashOperator
from datetime import datetime
from shared_utils.dag import DEFAULT_ARGS, DEFAULT_SETTINGS
from shared_utils.hash import merge_dict
import os


ABSOLUTE_PATH = os.path.abspath(os.path.dirname(__file__))
pipeline_path = os.path.join(ABSOLUTE_PATH, 'pipelines', 'pipeline_test')

dag = DAG(
    'test_mage_pipeline',
    default_args=merge_dict(DEFAULT_ARGS, {
        'start_date': datetime(2022, 7, 14),
    }),
    description='Test running mage pipeline.',
    schedule_interval='@once',
    **DEFAULT_SETTINGS,
)

task = BashOperator(
    dag=dag,
    task_id='run_pipeline',
    bash_command=f'mage run {pipeline_path} prod_pipeline',
)

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime
from mage_ai.data_preparation.models.pipeline import Pipeline
from shared_utils.dag import DEFAULT_ARGS, DEFAULT_SETTINGS
from shared_utils.hash import ignore_keys, merge_dict
import os


ABSOLUTE_PATH = os.path.abspath(os.path.dirname(__file__))
pipeline_path = os.path.join(ABSOLUTE_PATH, 'pipelines', 'pipeline_test')

dag = DAG(
    'test_mage_dag',
    default_args=merge_dict(DEFAULT_ARGS, {
        'start_date': datetime(2022, 7, 14),
    }),
    description='Test running mage pipeline.',
    schedule_interval='@once',
    **DEFAULT_SETTINGS,
)

pipeline = Pipeline('prod_pipeline', repo_path=pipeline_path)

blocks = pipeline.blocks_by_uuid

tasks = []


def build_execute_block(block):
    def _callable(ds, **kwargs):
        block.execute_sync(
            analyze_outputs=False,
            update_status=False,
        )

    return _callable


for uuid, b in blocks.items():
    if b.type == 'scratchpad':
        continue
    tasks.append(dict(
        task_id=uuid,
        upstream_task_ids=b.upstream_block_uuids,
        python_callable=build_execute_block(b),
    ))


def initialize_tasks(dag, tasks):
    operators = {}

    for task_dict in tasks:
        task_operator = PythonOperator(dag=dag, **ignore_keys(task_dict, [
            'upstream_task_ids',
        ]))
        operators[task_operator.task_id] = task_operator

    # Must set the dependencies after adding the task to the DAG or else error
    for task_dict in tasks:
        for task_id in task_dict.get('upstream_task_ids', []):
            dag.set_dependency(task_id, task_dict['task_id'])

    return operators


initialize_tasks(dag, tasks)