#!/bin/sh
# Simple Docker Database Query Tool

set -e

# Check docker is available and accessible
if ! docker ps >/dev/null 2>&1; then
    echo "[ERROR] Cannot access Docker daemon" >&2
    echo "[WARNING] Make sure you are in the 'docker' group or run: newgrp docker" >&2
    exit 1
fi

# Detect database type from image name
detect_db_type() {
    case "$1" in
        *postgres*|*postgresql*) echo "postgres" ;;
        *mysql*|*mariadb*) echo "mysql" ;;
        *mongo*) echo "mongo" ;;
        *redis*) echo "redis" ;;
        *) echo "unknown" ;;
    esac
}

# Parse drizzle config to find migration output directory
find_migration_dir_from_config() {
    config_file="$1"

    # Try to find 'out' field in config file
    # Supports both single and double quotes: out: './drizzle', out: "./migrations"
    out_dir=$(grep "out:" "$config_file" | \
              sed -E "s/.*['\"]([^'\"]+)['\"].*/\1/" | \
              head -1)

    if [ -n "$out_dir" ]; then
        echo "$out_dir"
        return 0
    fi
    return 1
}

# Find drizzle migration files
find_drizzle_migrations() {
    # Try to find git root first
    git_root=$(git rev-parse --show-toplevel 2>/dev/null)

    if [ -n "$git_root" ]; then
        search_base="$git_root"
        echo "[INFO] Found git root: $git_root" >&2
    else
        search_base="."
        echo "[INFO] Not in git repo, using current directory" >&2
    fi

    # Try to find drizzle config file
    for config_name in drizzle.config.ts drizzle.config.js drizzle.config.mjs drizzle.config.json; do
        config_path="$search_base/$config_name"
        if [ -f "$config_path" ]; then
            echo "[INFO] Found config: $config_name" >&2
            migration_dir=$(find_migration_dir_from_config "$config_path")

            if [ -n "$migration_dir" ]; then
                # Remove leading ./ if present
                migration_dir=$(echo "$migration_dir" | sed 's|^\./||')
                full_path="$search_base/$migration_dir"

                echo "[INFO] Using migration directory from config: $migration_dir" >&2

                if [ -d "$full_path" ]; then
                    find "$full_path" -name "*.sql" -type f 2>/dev/null | sort -r
                    return 0
                else
                    echo "[WARNING] Config specifies '$migration_dir' but directory not found" >&2
                fi
            fi
        fi
    done

    echo "[INFO] No drizzle config found, searching default locations..." >&2

    # Fallback: Search for common drizzle migration folders
    for dir in drizzle migrations db/migrations src/db/migrations; do
        full_path="$search_base/$dir"
        if [ -d "$full_path" ]; then
            find "$full_path" -name "*.sql" -type f 2>/dev/null | sort -r
            return 0
        fi
    done
    return 1
}

# Parse command line arguments
MIGRATE_MODE=0
GENERATE_MODE=0

if [ "$1" = "-h" ] || [ "$1" = "--help" ]; then
    echo "Usage: $(basename "$0") [OPTIONS]"
    echo ""
    echo "Interactive script to query Docker database containers"
    echo ""
    echo "Options:"
    echo "  -h, --help              Show this help message"
    echo "  -g, --generate          Run drizzle-kit generate to create migration files"
    echo "  -m, --migrate           Run drizzle migration SQL file"
    echo "  -gm, --generate-migrate Run drizzle-kit generate then immediately migrate"
    echo ""
    echo "Supported databases:"
    echo "  - PostgreSQL"
    echo "  - MySQL/MariaDB"
    echo "  - MongoDB"
    echo "  - Redis"
    exit 0
elif [ "$1" = "-g" ] || [ "$1" = "--generate" ]; then
    GENERATE_MODE=1
elif [ "$1" = "-m" ] || [ "$1" = "--migrate" ]; then
    MIGRATE_MODE=1
elif [ "$1" = "-gm" ] || [ "$1" = "--generate-migrate" ]; then
    GENERATE_MODE=1
    MIGRATE_MODE=1
fi

# Handle generate step (runs before container selection)
if [ "$GENERATE_MODE" = 1 ]; then
    git_root=$(git rev-parse --show-toplevel 2>/dev/null)
    if [ -n "$git_root" ]; then
        project_dir="$git_root"
        echo "[INFO] Found git root: $git_root"
    else
        project_dir="$(pwd)"
        echo "[INFO] Using current directory: $project_dir"
    fi

    echo "[INFO] Running drizzle-kit generate..."
    echo ""

    if (cd "$project_dir" && npx drizzle-kit generate); then
        echo ""
        echo "[SUCCESS] Migration files generated successfully"
    else
        echo ""
        echo "[ERROR] drizzle-kit generate failed" >&2
        exit 1
    fi

    if [ "$MIGRATE_MODE" = 0 ]; then
        exit 0
    fi

    echo ""
fi

# Get database containers and save to temp file
TMP_FILE="/tmp/ddb_$$.tmp"
trap 'rm -f "$TMP_FILE"' EXIT

docker ps --format "{{.ID}}|{{.Names}}|{{.Image}}|{{.Status}}" | \
while IFS='|' read -r id name image status; do
    db_type=$(detect_db_type "$image")
    if [ "$db_type" != "unknown" ]; then
        echo "$id|$name|$image|$db_type|$status"
    fi
done > "$TMP_FILE"

# Check if any database containers found
if [ ! -s "$TMP_FILE" ]; then
    echo "[ERROR] No running database containers found" >&2
    exit 1
fi

# Display menu
echo ""
echo "=== Running Database Containers ==="
echo ""

i=1
while IFS='|' read -r id name image db_type status; do
    echo "  $i) $name [$db_type] - $id"
    i=$((i + 1))
done < "$TMP_FILE"

# Get selection
echo ""
total=$(wc -l < "$TMP_FILE")
if [ "$total" -eq 1 ]; then
    echo -n "Select container [1]: "
    read selection
    selection="${selection:-1}"
else
    echo -n "Select container (1-$total): "
    read selection
fi

# Validate selection
case "$selection" in
    ''|*[!0-9]*)
        echo "[ERROR] Invalid selection" >&2
        exit 1
        ;;
esac

if [ "$selection" -lt 1 ] || [ "$selection" -gt "$total" ]; then
    echo "[ERROR] Selection out of range" >&2
    exit 1
fi

# Get selected container info
container_info=$(sed -n "${selection}p" "$TMP_FILE")
container_id=$(echo "$container_info" | cut -d'|' -f1)
container_name=$(echo "$container_info" | cut -d'|' -f2)
db_type=$(echo "$container_info" | cut -d'|' -f4)

echo ""
echo "[INFO] Selected: $container_name ($db_type)"
echo ""

# Handle migrate mode or interactive query mode
if [ "$MIGRATE_MODE" = 1 ]; then
    # Migrate mode - find and select SQL file
    echo "[INFO] Searching for drizzle migration files..."
    echo ""

    SQL_FILES_TMP="/tmp/ddb_sql_files_$$.tmp"
    trap 'rm -f "$TMP_FILE" "$SQL_FILES_TMP"' EXIT

    if ! find_drizzle_migrations > "$SQL_FILES_TMP"; then
        echo "[ERROR] No drizzle migration folder found" >&2
        echo "[INFO] Searched in: drizzle/, migrations/, db/migrations/, src/db/migrations/" >&2
        exit 1
    fi

    if [ ! -s "$SQL_FILES_TMP" ]; then
        echo "[ERROR] No SQL files found in migration folders" >&2
        exit 1
    fi

    echo "=== Available Migration Files ==="
    echo ""

    i=1
    while IFS= read -r sql_file; do
        echo "  $i) $(basename "$sql_file")"
        i=$((i + 1))
    done < "$SQL_FILES_TMP"

    echo ""
    sql_total=$(wc -l < "$SQL_FILES_TMP")
    echo -n "Select SQL file(s) to execute (e.g. 3 or 3,5,7 or 3-7) [1-$sql_total]: "
    read sql_selection

    if [ -z "$sql_selection" ]; then
        echo "[ERROR] No selection provided" >&2
        exit 1
    fi

    # Parse selection into reverse-sorted list (highest number first = oldest file first)
    SELECTED_TMP="/tmp/ddb_selected_$$.tmp"
    trap 'rm -f "$TMP_FILE" "$SQL_FILES_TMP" "$SELECTED_TMP"' EXIT

    # Split by comma and expand ranges
    echo "$sql_selection" | tr ',' '\n' | while IFS= read -r part; do
        part=$(echo "$part" | tr -d ' ')
        case "$part" in
            *-*)
                range_start=$(echo "$part" | cut -d'-' -f1)
                range_end=$(echo "$part" | cut -d'-' -f2)
                # Validate range parts are numbers
                case "$range_start" in ''|*[!0-9]*) echo "[ERROR] Invalid range: $part" >&2; exit 1;; esac
                case "$range_end" in ''|*[!0-9]*) echo "[ERROR] Invalid range: $part" >&2; exit 1;; esac
                n=$range_start
                while [ "$n" -le "$range_end" ]; do
                    echo "$n"
                    n=$((n + 1))
                done
                ;;
            *)
                case "$part" in ''|*[!0-9]*) echo "[ERROR] Invalid selection: $part" >&2; exit 1;; esac
                echo "$part"
                ;;
        esac
    done | sort -rn -u > "$SELECTED_TMP"

    if [ ! -s "$SELECTED_TMP" ]; then
        echo "[ERROR] Invalid selection" >&2
        exit 1
    fi

    # Validate all selections are in range
    while IFS= read -r num; do
        if [ "$num" -lt 1 ] || [ "$num" -gt "$sql_total" ]; then
            echo "[ERROR] Selection $num out of range (1-$sql_total)" >&2
            exit 1
        fi
    done < "$SELECTED_TMP"

    selected_count=$(wc -l < "$SELECTED_TMP")

    # Show selected files (in execution order: ascending = oldest first)
    echo ""
    echo "=== Selected Migration Files (execution order) ==="
    echo ""
    while IFS= read -r num; do
        sql_file_path=$(sed -n "${num}p" "$SQL_FILES_TMP")
        echo "  $num) $(basename "$sql_file_path")"
    done < "$SELECTED_TMP"

    # Preview if single file
    if [ "$selected_count" -eq 1 ]; then
        sql_file_path=$(sed -n "$(cat "$SELECTED_TMP")p" "$SQL_FILES_TMP")
        echo ""
        echo "=== SQL Content Preview ==="
        head -20 "$sql_file_path"

        if [ "$(wc -l < "$sql_file_path")" -gt 20 ]; then
            echo "..."
            echo "(showing first 20 lines)"
        fi
    fi

    echo ""
    if [ "$selected_count" -gt 1 ]; then
        echo -n "Execute $selected_count migrations? (y/N): "
    else
        echo -n "Execute this migration? (y/N): "
    fi
    read confirm

    case "$confirm" in
        y|Y|yes|YES)
            # Build per-file SQL temps for sequential execution
            QUERY_TMP_DIR="/tmp/ddb_queries_$$"
            mkdir -p "$QUERY_TMP_DIR"
            trap 'rm -f "$TMP_FILE" "$SQL_FILES_TMP" "$SELECTED_TMP"; rm -rf "$QUERY_TMP_DIR"' EXIT

            migrate_seq=0
            while IFS= read -r num; do
                sql_file_path=$(sed -n "${num}p" "$SQL_FILES_TMP")
                migrate_seq=$((migrate_seq + 1))
                { printf 'BEGIN;\n'; cat "$sql_file_path"; printf '\nCOMMIT;\n'; } > "$QUERY_TMP_DIR/${migrate_seq}_$(basename "$sql_file_path")"
            done < "$SELECTED_TMP"

            query="__MIGRATE__"
            ;;
        *)
            echo "[INFO] Migration cancelled"
            exit 0
            ;;
    esac
else
    # Normal mode - ask for query or interactive shell
    echo -n "Enter query (or press Enter for interactive shell): "
    read query
fi

if [ -z "$query" ]; then
    # Interactive shell
    echo "[INFO] Starting interactive shell for $container_name ($db_type)..."
    echo ""

    case "$db_type" in
        postgres)
            db_name=$(docker exec "$container_id" sh -c 'echo $POSTGRES_DB' 2>/dev/null || echo "postgres")
            user=$(docker exec "$container_id" sh -c 'echo $POSTGRES_USER' 2>/dev/null || echo "postgres")
            docker exec -it "$container_id" psql -U "$user" -d "$db_name"
            ;;
        mysql)
            db_name=$(docker exec "$container_id" sh -c 'echo $MYSQL_DATABASE' 2>/dev/null || echo "")
            user=$(docker exec "$container_id" sh -c 'echo $MYSQL_USER' 2>/dev/null || echo "root")
            password=$(docker exec "$container_id" sh -c 'echo $MYSQL_ROOT_PASSWORD' 2>/dev/null || echo "")

            if [ -n "$password" ]; then
                docker exec -it "$container_id" mysql -u"$user" -p"$password" ${db_name:+-D "$db_name"}
            else
                docker exec -it "$container_id" mysql -u"$user" ${db_name:+-D "$db_name"}
            fi
            ;;
        mongo)
            db_name=$(docker exec "$container_id" sh -c 'echo $MONGO_INITDB_DATABASE' 2>/dev/null || echo "test")
            docker exec -it "$container_id" mongosh "$db_name"
            ;;
        redis)
            docker exec -it "$container_id" redis-cli
            ;;
        *)
            echo "[ERROR] Unsupported database type: $db_type" >&2
            exit 1
            ;;
    esac
else
    # Execute query
    echo "[INFO] Executing query on $container_name ($db_type)..."
    echo ""

    case "$db_type" in
        postgres)
            db_name=$(docker exec "$container_id" sh -c 'echo $POSTGRES_DB' 2>/dev/null || echo "postgres")
            user=$(docker exec "$container_id" sh -c 'echo $POSTGRES_USER' 2>/dev/null || echo "postgres")

            if [ "$MIGRATE_MODE" = 1 ]; then
                migrate_failed=0
                for qfile in $(ls "$QUERY_TMP_DIR"/*.sql | sort -n); do
                    fname=$(basename "$qfile" | sed 's/^[0-9]*_//')
                    echo "[INFO] Executing: $fname"
                    docker exec -i "$container_id" psql -v ON_ERROR_STOP=1 -U "$user" -d "$db_name" < "$qfile"
                    if [ $? -ne 0 ]; then
                        echo ""
                        echo "[ERROR] Migration failed: $fname" >&2
                        migrate_failed=1
                        break
                    fi
                    echo "[SUCCESS] Applied: $fname"
                    echo ""
                done
                if [ "$migrate_failed" -eq 1 ]; then
                    exit 1
                fi
                echo "[SUCCESS] All migrations executed successfully"
            else
                docker exec -it "$container_id" psql -U "$user" -d "$db_name" -c "$query"
                psql_exit=$?

                if [ "$psql_exit" -eq 0 ]; then
                    echo ""
                    echo "[SUCCESS] Query executed successfully"
                else
                    echo ""
                    echo "[ERROR] Query execution failed" >&2
                    exit 1
                fi
            fi
            ;;
        mysql)
            db_name=$(docker exec "$container_id" sh -c 'echo $MYSQL_DATABASE' 2>/dev/null || echo "")
            user=$(docker exec "$container_id" sh -c 'echo $MYSQL_USER' 2>/dev/null || echo "root")
            password=$(docker exec "$container_id" sh -c 'echo $MYSQL_ROOT_PASSWORD' 2>/dev/null || echo "")

            if [ "$MIGRATE_MODE" = 1 ]; then
                migrate_failed=0
                for qfile in $(ls "$QUERY_TMP_DIR"/*.sql | sort -n); do
                    fname=$(basename "$qfile" | sed 's/^[0-9]*_//')
                    echo "[INFO] Executing: $fname"
                    if [ -n "$password" ]; then
                        docker exec -i "$container_id" mysql -u"$user" -p"$password" ${db_name:+-D "$db_name"} < "$qfile"
                    else
                        docker exec -i "$container_id" mysql -u"$user" ${db_name:+-D "$db_name"} < "$qfile"
                    fi
                    if [ $? -ne 0 ]; then
                        echo ""
                        echo "[ERROR] Migration failed: $fname" >&2
                        migrate_failed=1
                        break
                    fi
                    echo "[SUCCESS] Applied: $fname"
                    echo ""
                done
                if [ "$migrate_failed" -eq 1 ]; then
                    exit 1
                fi
                echo "[SUCCESS] All migrations executed successfully"
            else
                if [ -n "$password" ]; then
                    docker exec -it "$container_id" mysql -u"$user" -p"$password" ${db_name:+-D "$db_name"} -e "$query"
                else
                    docker exec -it "$container_id" mysql -u"$user" ${db_name:+-D "$db_name"} -e "$query"
                fi

                if [ $? -eq 0 ]; then
                    echo ""
                    echo "[SUCCESS] Query executed successfully"
                else
                    echo ""
                    echo "[ERROR] Query execution failed" >&2
                    exit 1
                fi
            fi
            ;;
        mongo)
            db_name=$(docker exec "$container_id" sh -c 'echo $MONGO_INITDB_DATABASE' 2>/dev/null || echo "test")

            if [ "$MIGRATE_MODE" = 1 ]; then
                migrate_failed=0
                for qfile in $(ls "$QUERY_TMP_DIR"/*.sql | sort -n); do
                    fname=$(basename "$qfile" | sed 's/^[0-9]*_//')
                    echo "[INFO] Executing: $fname"
                    docker exec -i "$container_id" mongosh "$db_name" < "$qfile"
                    if [ $? -ne 0 ]; then
                        echo ""
                        echo "[ERROR] Migration failed: $fname" >&2
                        migrate_failed=1
                        break
                    fi
                    echo "[SUCCESS] Applied: $fname"
                    echo ""
                done
                if [ "$migrate_failed" -eq 1 ]; then
                    exit 1
                fi
                echo "[SUCCESS] All migrations executed successfully"
            else
                docker exec -it "$container_id" mongosh "$db_name" --eval "$query"

                if [ $? -eq 0 ]; then
                    echo ""
                    echo "[SUCCESS] Query executed successfully"
                else
                    echo ""
                    echo "[ERROR] Query execution failed" >&2
                    exit 1
                fi
            fi
            ;;
        redis)
            if docker exec -it "$container_id" redis-cli "$query"; then
                echo ""
                echo "[SUCCESS] Command executed successfully"
            else
                echo ""
                echo "[ERROR] Command execution failed" >&2
                exit 1
            fi
            ;;
        *)
            echo "[ERROR] Unsupported database type: $db_type" >&2
            exit 1
            ;;
    esac
fi