"""Per-database light data generation DAGs.

Each DAG runs one generator script and accepts a `rows` value via the
dag_run conf (passed by the Command Center), exported as GEN_ROWS.
Triggerable independently so every database gets its own button.
"""
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime
import os
import subprocess

SCRIPTS_DIR = "/opt/airflow/dags/scripts"

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

DEFAULT_ROWS = {"neo4j": "2000"}

default_args = {"owner": "airflow", "retries": 0}


def make_runner(script: str, default_rows: str):
    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


for _src, _script in SOURCES.items():
    _dag_id = f"gen_{_src}"
    _dag = DAG(
        dag_id=_dag_id,
        default_args=default_args,
        description=f"Generate light data into {_src} (GEN_ROWS via conf.rows)",
        schedule=None,
        start_date=datetime(2025, 1, 1),
        catchup=False,
        tags=["data", "generation", _src],
    )
    PythonOperator(
        task_id=f"generate_{_src}",
        python_callable=make_runner(_script, DEFAULT_ROWS.get(_src, "5000")),
        dag=_dag,
    )
    globals()[_dag_id] = _dag