#!/bin/bash
set -e

# Activate uv virtual environment if it exists
if [ -f "/opt/airflow/.venv/bin/activate" ]; then
    echo "Activating uv virtual environment..."
    source /opt/airflow/.venv/bin/activate
fi

# Wait for postgres to be ready
wait_for_postgres() {
    # Extract host and port from DATABASE_URL (format: postgresql+asyncpg://user:pass@host:port/db)
    local db_host db_port host_port
    host_port="${DATABASE_URL#*@}"   # strip everything before @
    host_port="${host_port%%/*}"     # strip everything after first /
    db_host="${host_port%%:*}"       # part before :
    db_port="${host_port##*:}"       # part after :
    db_host="${db_host:-localhost}"
    db_port="${db_port:-5432}"

    echo "Waiting for PostgreSQL at ${db_host}:${db_port}..."
    while ! nc -z "${db_host}" "${db_port}"; do
        sleep 1
    done
    echo "PostgreSQL is ready!"
}

# Create admin user if FABAuthManager is configured
create_www_user() {
    local local_password=""

    # Get password from env var
    if [[ -n "${_AIRFLOW_WWW_USER_PASSWORD_CMD:-}" ]]; then
        local_password=$(eval "${_AIRFLOW_WWW_USER_PASSWORD_CMD}")
    elif [[ -n "${_AIRFLOW_WWW_USER_PASSWORD:-}" ]]; then
        local_password="${_AIRFLOW_WWW_USER_PASSWORD}"
    fi

    if [[ -z "${local_password}" ]]; then
        echo "No _AIRFLOW_WWW_USER_PASSWORD set, skipping user creation"
        return
    fi

    # Check if FABAuthManager is configured
    if airflow config get-value core auth_manager 2>/dev/null | grep -q "FabAuthManager"; then
        echo "Creating admin user with FABAuthManager..."
        airflow users create \
            --username "${_AIRFLOW_WWW_USER_USERNAME:-admin}" \
            --firstname "${_AIRFLOW_WWW_USER_FIRSTNAME:-Airflow}" \
            --lastname "${_AIRFLOW_WWW_USER_LASTNAME:-Admin}" \
            --email "${_AIRFLOW_WWW_USER_EMAIL:-airflowadmin@example.com}" \
            --role "${_AIRFLOW_WWW_USER_ROLE:-Admin}" \
            --password "${local_password}" || true
        echo "Admin user created successfully!"
    else
        echo "Auth manager is not FABAuthManager, skipping user creation"
        echo "SimpleAuthManager credentials will be auto-generated"
    fi
}

# Initialize Airflow if needed
init_airflow() {
    echo "Running Airflow database migrations..."
    airflow db migrate
    echo "Airflow database migration done!"

    # Run application schema migrations (Alembic)
    echo "Running application schema migrations (Alembic)..."
    cd /opt/airflow
    /opt/airflow/.venv/bin/alembic upgrade head || {
        echo "Warning: Alembic migrations failed. This might be expected on first run."
    }
    echo "Application schema migration done!"

    # Create admin user if _AIRFLOW_WWW_USER_CREATE is set
    if [[ "${_AIRFLOW_WWW_USER_CREATE:-}" == "true" ]]; then
        create_www_user
    fi
}

# Main
wait_for_postgres
init_airflow

# Execute the main command (prepend 'airflow' if needed)
# Note: Airflow 3.x renamed 'webserver' to 'api-server'
if [ "$1" = "webserver" ]; then
    # Translate webserver to api-server for Airflow 3.x
    shift
    exec airflow api-server "$@"
elif [ "$1" = "scheduler" ] || [ "$1" = "worker" ] || [ "$1" = "triggerer" ] || [ "$1" = "api-server" ] || [ "$1" = "dag-processor" ]; then
    exec airflow "$@"
else
    exec "$@"
fi