#!/bin/bash
set -eE

# Global variables set by parse_args
MIGRATION_TYPE=""
FRESH_MODE=false

parse_args() {
    MIGRATION_TYPE="$1"
    shift || true

    for arg in "$@"; do
        case "$arg" in
            --fresh)
                FRESH_MODE=true
                ;;
            *)
                echo "Unknown option: $arg"
                exit 1
                ;;
        esac
    done
}

show_usage() {
    echo "Usage: $0 <all|persons|cyclotron|cyclotron-node|behavioral-cohorts|flags-read-store> [--fresh]"
    echo "  all                - Run all migrations"
    echo "  persons            - Run only persons migrations"
    echo "  cyclotron          - Run only cyclotron migrations"
    echo "  cyclotron-node     - Run only cyclotron node migrations"
    echo "  behavioral-cohorts - Run only behavioral cohorts migrations"
    echo "  flags-read-store   - Run only flags read store migrations"
    echo ""
    echo "Options:"
    echo "  --fresh            - Drop and recreate database before migrating (for test environments)"
}

parse_args "$@"

log_json() {
    local level="$1"
    local message="$2"
    local status="$3"
    local migration_type="${4:-$MIGRATION_TYPE}"
    local extra="${5:-}"

    local timestamp
    timestamp=$(date -u +"%Y-%m-%dT%H:%M:%SZ")

    local json="{\"timestamp\":\"$timestamp\",\"level\":\"$level\",\"message\":\"$message\",\"migration_type\":\"$migration_type\",\"status\":\"$status\""

    if [ -n "$extra" ]; then
        json="$json,$extra"
    fi

    json="$json}"
    echo "$json"
}

handle_error() {
    local exit_code=$?
    log_json "error" "Migration failed" "failed" "$MIGRATION_TYPE" "\"exit_code\":$exit_code"
    exit $exit_code
}

# Trap errors and use the error handler to log the error as JSON
trap handle_error ERR

if [ -z "$MIGRATION_TYPE" ]; then
    log_json "error" "No migration type specified" "failed" "unknown"
    echo ""
    show_usage
    exit 1
fi

# Determine the base directory for migrations
# In Docker, we use /migrations; locally, we use the parent directory
if [ -d "/migrations" ]; then
    MIGRATIONS_BASE="/migrations"
else
    SCRIPT_DIR=$(dirname "$(readlink -f "$0")")
    MIGRATIONS_BASE="$SCRIPT_DIR/.."
fi

run_persons_migrations() {
    local db_type="persons"
    PERSONS_DATABASE_NAME=${PERSONS_DATABASE_NAME:-posthog_persons}
    PERSONS_DATABASE_URL=${PERSONS_DATABASE_URL:-postgres://posthog:posthog@localhost:5432/$PERSONS_DATABASE_NAME}

    log_json "info" "Starting migrations" "starting" "$db_type" "\"database_name\":\"$PERSONS_DATABASE_NAME\",\"fresh_mode\":$FRESH_MODE"

    if [ "$FRESH_MODE" = true ]; then
        sqlx database reset -y -D "$PERSONS_DATABASE_URL" --source "$MIGRATIONS_BASE/persons_migrations/"
        log_json "info" "Database reset and migrations completed" "completed" "$db_type"
    else
        sqlx database create -D "$PERSONS_DATABASE_URL"
        log_json "info" "Database created or already exists" "running" "$db_type"

        sqlx migrate run -D "$PERSONS_DATABASE_URL" --source "$MIGRATIONS_BASE/persons_migrations/"
        log_json "info" "Migrations completed successfully" "completed" "$db_type"
    fi
}

run_cyclotron_migrations() {
    local db_type="cyclotron"
    CYCLOTRON_DATABASE_NAME=${CYCLOTRON_DATABASE_NAME:-cyclotron}
    CYCLOTRON_DATABASE_URL=${CYCLOTRON_DATABASE_URL:-postgres://posthog:posthog@localhost:5432/$CYCLOTRON_DATABASE_NAME}

    log_json "info" "Starting migrations" "starting" "$db_type" "\"database_name\":\"$CYCLOTRON_DATABASE_NAME\",\"fresh_mode\":$FRESH_MODE"

    if [ "$FRESH_MODE" = true ]; then
        sqlx database reset -y -D "$CYCLOTRON_DATABASE_URL" --source "$MIGRATIONS_BASE/cyclotron-core/migrations/"
        log_json "info" "Database reset and migrations completed" "completed" "$db_type"
    else
        sqlx database create -D "$CYCLOTRON_DATABASE_URL"
        log_json "info" "Database created or already exists" "running" "$db_type"

        sqlx migrate run -D "$CYCLOTRON_DATABASE_URL" --source "$MIGRATIONS_BASE/cyclotron-core/migrations/"
        log_json "info" "Migrations completed successfully" "completed" "$db_type"
    fi
}

run_cyclotron_node_migrations() {
    local db_type="cyclotron-node"
    CYCLOTRON_NODE_DATABASE_NAME=${CYCLOTRON_NODE_DATABASE_NAME:-cyclotron_node}
    CYCLOTRON_NODE_DATABASE_URL=${CYCLOTRON_NODE_DATABASE_URL:-postgres://posthog:posthog@localhost:5432/$CYCLOTRON_NODE_DATABASE_NAME}

    log_json "info" "Starting migrations" "starting" "$db_type" "\"database_name\":\"$CYCLOTRON_NODE_DATABASE_NAME\",\"fresh_mode\":$FRESH_MODE"

    if [ "$FRESH_MODE" = true ]; then
        sqlx database reset -y -D "$CYCLOTRON_NODE_DATABASE_URL" --source "$MIGRATIONS_BASE/cyclotron-node-migrations/"
        log_json "info" "Database reset and migrations completed" "completed" "$db_type"
    else
        sqlx database create -D "$CYCLOTRON_NODE_DATABASE_URL"
        log_json "info" "Database created or already exists" "running" "$db_type"

        sqlx migrate run -D "$CYCLOTRON_NODE_DATABASE_URL" --source "$MIGRATIONS_BASE/cyclotron-node-migrations/"
        log_json "info" "Migrations completed successfully" "completed" "$db_type"
    fi
}

run_behavioral_cohorts_migrations() {
    local db_type="behavioral-cohorts"
    BEHAVIORAL_COHORTS_DATABASE_NAME=${BEHAVIORAL_COHORTS_DATABASE_NAME:-behavioral_cohorts}

    # Use environment variables with defaults
    DB_HOST="${POSTGRES_BEHAVIORAL_COHORTS_HOST:-localhost}"
    DB_USER="${POSTGRES_BEHAVIORAL_COHORTS_USER:-posthog}"
    DB_PASS="${POSTGRES_BEHAVIORAL_COHORTS_PASSWORD:-posthog}"

    BEHAVIORAL_COHORTS_DATABASE_URL=${BEHAVIORAL_COHORTS_DATABASE_URL:-postgres://${DB_USER}:${DB_PASS}@${DB_HOST}:5432/$BEHAVIORAL_COHORTS_DATABASE_NAME}

    log_json "info" "Starting migrations" "starting" "$db_type" "\"database_name\":\"$BEHAVIORAL_COHORTS_DATABASE_NAME\",\"fresh_mode\":$FRESH_MODE"

    if [ "$FRESH_MODE" = true ]; then
        sqlx database reset -y -D "$BEHAVIORAL_COHORTS_DATABASE_URL" --source "$MIGRATIONS_BASE/behavioral_cohorts_migrations/"
        log_json "info" "Database reset and migrations completed" "completed" "$db_type"
    else
        sqlx database create -D "$BEHAVIORAL_COHORTS_DATABASE_URL"
        log_json "info" "Database created or already exists" "running" "$db_type"

        sqlx migrate run -D "$BEHAVIORAL_COHORTS_DATABASE_URL" --source "$MIGRATIONS_BASE/behavioral_cohorts_migrations/"
        log_json "info" "Migrations completed successfully" "completed" "$db_type"
    fi
}

run_flags_read_store_migrations() {
    local db_type="flags-read-store"
    FLAGS_READ_STORE_DATABASE_NAME=${FLAGS_READ_STORE_DATABASE_NAME:-flags_read_store}

    # Use environment variables with defaults
    DB_HOST="${POSTGRES_FLAGS_READ_STORE_HOST:-localhost}"
    DB_USER="${POSTGRES_FLAGS_READ_STORE_USER:-posthog}"
    DB_PASS="${POSTGRES_FLAGS_READ_STORE_PASSWORD:-posthog}"

    FLAGS_READ_STORE_DATABASE_URL=${FLAGS_READ_STORE_DATABASE_URL:-postgres://${DB_USER}:${DB_PASS}@${DB_HOST}:5432/$FLAGS_READ_STORE_DATABASE_NAME}

    log_json "info" "Starting migrations" "starting" "$db_type" "\"database_name\":\"$FLAGS_READ_STORE_DATABASE_NAME\",\"fresh_mode\":$FRESH_MODE"

    if [ "$FRESH_MODE" = true ]; then
        sqlx database reset -y -D "$FLAGS_READ_STORE_DATABASE_URL" --source "$MIGRATIONS_BASE/flags_read_store_migrations/"
        log_json "info" "Database reset and migrations completed" "completed" "$db_type"
    else
        sqlx database create -D "$FLAGS_READ_STORE_DATABASE_URL"
        log_json "info" "Database created or already exists" "running" "$db_type"

        sqlx migrate run -D "$FLAGS_READ_STORE_DATABASE_URL" --source "$MIGRATIONS_BASE/flags_read_store_migrations/"
        log_json "info" "Migrations completed successfully" "completed" "$db_type"
    fi
}

case "$MIGRATION_TYPE" in
    persons)
        run_persons_migrations
        ;;
    cyclotron)
        run_cyclotron_migrations
        ;;
    cyclotron-node)
        run_cyclotron_node_migrations
        ;;
    behavioral-cohorts)
        run_behavioral_cohorts_migrations
        ;;
    flags-read-store)
        run_flags_read_store_migrations
        ;;
    all)
        run_persons_migrations
        run_cyclotron_migrations
        run_cyclotron_node_migrations
        run_behavioral_cohorts_migrations
        ;;
    *)
        log_json "error" "Invalid migration type specified" "failed" "$MIGRATION_TYPE"
        show_usage
        exit 1
        ;;
esac

log_json "info" "All requested migrations completed successfully" "completed"
