"""Generate fake data for all databases in one run.

Runs all five generator scripts in parallel. Honors a `rows` value passed via
the dag_run conf (from the Command Center), exported to each script as
GEN_ROWS so the "All sources" button respects the row count.
"""
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime, timedelta
import os
import subprocess

default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2025, 1, 1),
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=2),
}

dag = DAG(
    'generate_data_all_databases',
    default_args=default_args,
    description='Generate fake data for all databases',
    schedule=None,
    catchup=False,
    tags=['data', 'generation', 'fake-data'],
)

SCRIPTS_DIR = '/opt/airflow/dags/scripts'

SCRIPTS = {
    'postgres': 'generate_postgres_sales_data.py',
    'mongodb': 'generate_mongodb_events_data.py',
    'mysql': 'generate_mysql_employee_data.py',
    'neo4j': 'generate_neo4j_graph_data.py',
    'cassandra': 'generate_cassandra_telemetry_data.py',
}

DEFAULT_ROWS = {'neo4j': '2000'}


def make_runner(script, default_rows):
    def _run(**context):
        dag_run = context.get('dag_run')
        conf = (dag_run.conf if dag_run else {}) or {}
        rows = str(conf.get('rows') or default_rows)
        env = dict(os.environ)
        env['GEN_ROWS'] = rows
        print(f"Running {script} with GEN_ROWS={rows}")
        result = subprocess.run(
            ['python3', os.path.join(SCRIPTS_DIR, script)],
            capture_output=True, text=True, env=env,
        )
        if result.stdout:
            print(result.stdout[-4000:])
        if result.stderr:
            print("STDERR:", result.stderr[-4000:])
        if result.returncode != 0:
            raise Exception(f"{script} failed with return code {result.returncode}")
    return _run


tasks = []
for _src, _script in SCRIPTS.items():
    tasks.append(PythonOperator(
        task_id=f"generate_{_src}_data",
        python_callable=make_runner(_script, DEFAULT_ROWS.get(_src, '5000')),
        dag=dag,
    ))

# all in parallel
tasks