#!/bin/bash
set -euo pipefail

# ---- Helpers ----
log() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*"; }

# Graceful shutdown: forward signals to child processes
trap 'log "Caught signal, shutting down..."; kill -TERM $(jobs -p) 2>/dev/null; exit 0' SIGTERM SIGINT SIGQUIT

# ---- Load Configuration ----
log "Loading configuration..."
source /config.sh
validate_config

log "Configuration:"
log "  DB backend : $(echo "${AIRFLOW__DATABASE__SQL_ALCHEMY_CONN}" | sed -E 's|://[^@]*@|://***:***@|')"
log "  DAG repo   : ${DAG_REPO_URL} (branch: ${DAG_REPO_BRANCH})"
log "  Executor   : ${AIRFLOW__CORE__EXECUTOR}"
log "  Ports      : nginx=${NGINX_PORT}  webserver=${AIRFLOW_WEB_PORT}  refresh=${REFRESH_SERVICE_PORT}"

# ---- Ensure directories ----
mkdir -p "${DAGS_DIR}" "${LOGS_DIR}"

# ---- Wait for Database (if PostgreSQL) ----
if [[ "${AIRFLOW__DATABASE__SQL_ALCHEMY_CONN}" == postgresql* ]]; then
    DB_WAIT_TIMEOUT="${DB_WAIT_TIMEOUT:-60}"
    log "Waiting for PostgreSQL (timeout: ${DB_WAIT_TIMEOUT}s)..."
    elapsed=0
    while [ "${elapsed}" -lt "${DB_WAIT_TIMEOUT}" ]; do
        if python3 -c "
import sqlalchemy, sys
try:
    engine = sqlalchemy.create_engine('${AIRFLOW__DATABASE__SQL_ALCHEMY_CONN}', connect_args={'connect_timeout': 5})
    with engine.connect() as conn:
        conn.execute(sqlalchemy.text('SELECT 1'))
    sys.exit(0)
except Exception:
    sys.exit(1)
" 2>/dev/null; then
            log "PostgreSQL is ready."
            break
        fi
        sleep 3
        elapsed=$((elapsed + 3))
    done
    if [ "${elapsed}" -ge "${DB_WAIT_TIMEOUT}" ]; then
        log "[ERROR] PostgreSQL not reachable after ${DB_WAIT_TIMEOUT}s."
        log "Falling back to SQLite so the container still boots."
        export AIRFLOW__DATABASE__SQL_ALCHEMY_CONN="sqlite:///${AIRFLOW_HOME}/airflow.db"
        export AIRFLOW__CORE__EXECUTOR="SequentialExecutor"
    fi
fi

# ---- Build authenticated repo URL (if PAT is set) ----
GIT_CLONE_URL="${DAG_REPO_URL}"
if [ -n "${DAG_REPO_TOKEN}" ]; then
    # Inject token: https://github.com/... → https://<token>@github.com/...
    GIT_CLONE_URL=$(echo "${DAG_REPO_URL}" | sed "s|https://|https://${DAG_REPO_TOKEN}@|")
    log "Using authenticated URL for DAG repo (PAT token set)."
fi

# ---- Sync DAG Repository ----
_create_sample_dag() {
    log "Creating sample DAG for testing..."
    cat > "${DAGS_DIR}/sample_dag.py" << 'SAMPLE_EOF'
"""
Sample DAG — created automatically because no DAG repository was configured.
Replace this by setting DAG_REPO_URL or adding your own DAGs to /opt/airflow/dags.
"""
from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.bash import BashOperator

with DAG(
    dag_id="sample_hello_world",
    description="A sample DAG to verify Airflow is working",
    schedule=timedelta(hours=1),
    start_date=datetime(2024, 1, 1),
    catchup=False,
    tags=["sample", "test"],
) as dag:
    hello = BashOperator(
        task_id="say_hello",
        bash_command='echo "Hello from Airflow! $(date)"',
    )
    check_env = BashOperator(
        task_id="check_environment",
        bash_command='echo "Python: $(python3 --version)" && echo "Airflow Home: $AIRFLOW_HOME"',
    )
    hello >> check_env
SAMPLE_EOF
    log "Sample DAG created at ${DAGS_DIR}/sample_dag.py"
}

log "Syncing DAG repository..."
if [ ! -d "${DAGS_DIR}/.git" ]; then
    # Fresh clone — clean any stale files (except sample_dag) first
    rm -rf "${DAGS_DIR:?}"/* 2>/dev/null || true
    if git clone --depth 1 --branch "${DAG_REPO_BRANCH}" "${GIT_CLONE_URL}" "${DAGS_DIR}" 2>/dev/null; then
        log "DAG repository cloned."
    else
        log "[WARN] Could not clone DAG repo (URL unreachable or private without token)."
        _create_sample_dag
    fi
else
    cd "${DAGS_DIR}"
    # If PAT changed, update the remote URL
    git remote set-url origin "${GIT_CLONE_URL}" 2>/dev/null || true
    if git fetch origin "${DAG_REPO_BRANCH}" 2>/dev/null && git reset --hard "origin/${DAG_REPO_BRANCH}" 2>/dev/null; then
        log "DAG repository updated."
    else
        log "[WARN] Could not update DAG repo. Using existing DAGs."
    fi
fi

# ---- Database Migrations ----
log "Running database migrations..."
airflow db migrate 2>&1 | tail -5
log "Migrations complete."

# ---- Create Admin User (idempotent) ----
log "Ensuring admin user exists..."
if ! airflow users list 2>/dev/null | grep -q "${AIRFLOW_ADMIN_USER}"; then
    airflow users create \
        --username "${AIRFLOW_ADMIN_USER}" \
        --firstname Admin \
        --lastname User \
        --role Admin \
        --email "${AIRFLOW_ADMIN_USER}@example.com" \
        --password "${AIRFLOW_ADMIN_PASSWORD}"
    log "Admin user '${AIRFLOW_ADMIN_USER}' created."
else
    log "Admin user '${AIRFLOW_ADMIN_USER}' already exists — skipping."
fi

# ---- Start Supervisor ----
log "Starting services via supervisord..."
exec /usr/bin/supervisord -c /etc/supervisord.conf