#!/bin/bash

# Confirm with user before starting services
echo "This script:"
echo "  Will install all the dependencies you need to develop Suno."
echo ""
read -p "Continue (y/N) " confirm
if [[ $confirm != [yY] ]]; then
    echo "Aborting..."
    exit 1
fi

# Function to check if a command exists
check_command() {
    if ! command -v $1 &> /dev/null; then
        echo "Error: $1 is not installed"
        echo "Please install it using: $2"
        return 1
    fi
    return 0
}

# Function to check if nvm is available; separate from check_command because it's not a command
check_nvm() {
    if ! type nvm &> /dev/null; then
        echo "Error: nvm is not loaded in the current shell"
        echo "Please run: source $(brew --prefix nvm)/nvm.sh"
        return 1
    fi
    return 0
}

# Function to wait for PostgreSQL to be ready
wait_for_postgres() {
    echo "Waiting for PostgreSQL to be ready..."
    for i in {1..30}; do
        if psql -d template1 -c "SELECT 1" > /dev/null 2>&1; then
            echo "PostgreSQL is ready"
            return 0
        fi
        if [ $i -eq 30 ]; then
            echo "Error: PostgreSQL failed to become ready in time"
            return 1
        fi
        sleep 1
    done
}

# Check all required dependencies
echo "Checking required dependencies..."

# Check Homebrew
if ! check_command brew "Install Homebrew: /bin/bash -c \"\$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)\""; then
    echo "Error: Homebrew is not installed"
    echo "Installing Homebrew..."
    /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
    if [ $? -ne 0 ]; then
        echo "Error: Failed to install Homebrew"
        exit 1
    fi
    echo >> ~/.zprofile
    echo 'eval "$(/opt/homebrew/bin/brew shellenv)"' >> ~/.zprofile
    eval "$(/opt/homebrew/bin/brew shellenv)"
    echo "Homebrew installed successfully"
fi

# Check Git
if ! check_command git "brew install git"; then
    echo "Error: Git is not installed"
    echo "Installing Git..."
    brew install git
    if [ $? -ne 0 ]; then
        echo "Error: Failed to install Git"
        exit 1
    fi
    brew install --cask git-credential-manager
    if [ $? -ne 0 ]; then
        echo "Error: Failed to install Git Credential Manager"
        exit 1
    fi
    echo "Git installed successfully"
fi

# Check Python 3.10
if ! check_command python3.10 "brew install python@3.10"; then
    echo "Error: Python 3.10 is not installed"
    echo "Installing Python 3.10..."
    brew install python@3.10
    if [ $? -ne 0 ]; then
        echo "Error: Failed to install Python 3.10"
        exit 1
    fi
    alias python=/opt/homebrew/bin/python3.10
    echo "Python 3.10 installed successfully"
fi

# Check uv
if ! check_command uv "brew install uv"; then
    echo "Error: uv is not installed"
    echo "Installing uv..."
    brew install uv
    if [ $? -ne 0 ]; then
        echo "Error: Failed to install uv"
        exit 1
    fi
    echo "uv installed successfully"
fi

# UV seems to break brew shellenv, so we need to re-source it
eval "$(/opt/homebrew/bin/brew shellenv)"

# Check npm
if ! check_command npm "brew install npm"; then
    echo "Error: npm is not installed"
    echo "Installing npm..."
    brew install npm
    if [ $? -ne 0 ]; then
        echo "Error: Failed to install npm"
        exit 1
    fi
    echo "npm installed successfully"
fi

# Check nvm
if ! check_nvm; then
    echo "Error: nvm is not installed"
    echo "Installing nvm..."
    brew install nvm
    if [ $? -ne 0 ]; then
        echo "Error: Failed to install nvm"
        exit 1
    fi
    source $(brew --prefix nvm)/nvm.sh
    if [ $? -ne 0 ]; then
        echo "Error: Failed to source nvm"
        exit 1
    fi
    echo "nvm installed successfully"
fi

# Check package manager
if ! check_command pnpm "brew install pnpm"; then
    echo "Error: pnpm is not installed"
    echo "Installing pnpm..."
    brew install pnpm
    if [ $? -ne 0 ]; then
        echo "Error: Failed to install pnpm"
        exit 1
    fi
    echo "pnpm installed successfully"
fi


# Check lefthook
if ! check_command lefthook "brew install lefthook"; then
    echo "Error: lefthook is not installed"
    echo "Installing lefthook..."
    brew install lefthook
    if [ $? -ne 0 ]; then
        echo "Error: Failed to install lefthook"
        exit 1
    fi
    lefthook install
    echo "lefthook installed successfully"
fi

# # Check Docker
# if ! check_command docker "Install Docker Desktop from https://www.docker.com/products/docker-desktop"; then
#     echo "Go do that now. Run this script again when you're done."
#     exit 1
# fi

echo "All dependencies are installed!"
echo


# Check for .env file in ui/app-ui
if [ ! -f "./ui/app-ui/.env" ]; then
    echo "Error: .env file not found in ./ui/app-ui directory"
    echo "Please create a .env file with the required environment variables"
    echo "You can ask someone to send it to you!"
    exit 1
fi

# Check for AWS_PROFILE environment variable
if [ -z "$AWS_PROFILE" ]; then
    echo "Error: AWS_PROFILE environment variable is not set"
    echo "This is needed for accessing AWS services"
    exit 1
fi

# Check if any PostgreSQL service is running
if ! brew services list | grep -q "postgresql.*started"; then
    echo "No PostgreSQL service is running."

    # Check if any PostgreSQL version is installed
    if brew list | grep -q "^postgresql"; then
        echo "PostgreSQL is installed but not running. Starting PostgreSQL service..."
        # Prefer postgresql@17, then @16, @15, @14, then unversioned postgresql
        if brew list | grep -q "^postgresql@17$"; then
            PG_FORMULA="postgresql@17"
        elif brew list | grep -q "^postgresql@16$"; then
            PG_FORMULA="postgresql@16"
        elif brew list | grep -q "^postgresql@15$"; then
            PG_FORMULA="postgresql@15"
        elif brew list | grep -q "^postgresql@14$"; then
            PG_FORMULA="postgresql@14"
        else
            PG_FORMULA=$(brew list | grep "^postgresql" | head -n 1)
        fi
        echo "Starting $PG_FORMULA..."
        brew services start $PG_FORMULA
        if [ $? -ne 0 ]; then
            echo "Error: Failed to start PostgreSQL service"
            exit 1
        fi
        echo "PostgreSQL service started successfully"

        # Wait for PostgreSQL to be ready
        if ! wait_for_postgres; then
            exit 1
        fi
    else
        echo "No PostgreSQL installation found. Installing PostgreSQL 17..."
        brew install postgresql@17
        if [ $? -ne 0 ]; then
            echo "Error: Failed to install PostgreSQL 17"
            exit 1
        fi
        echo "PostgreSQL 17 installed successfully"
        brew services start postgresql@17
        if [ $? -ne 0 ]; then
            echo "Error: Failed to start PostgreSQL 17 service"
            exit 1
        fi
        echo "PostgreSQL 17 service started successfully"

        # Wait for PostgreSQL to be ready
        if ! wait_for_postgres; then
            exit 1
        fi
    fi
else
    echo "PostgreSQL service is already running"
fi

# Check if Redis service is running
if ! brew services list | grep -q "redis.*started"; then
    echo "Redis is not running. Starting Redis service..."
    brew services start redis
    if [ $? -ne 0 ]; then
        echo "Error: Failed to start Redis service. Did you install it?"
        echo "Installing Redis..."
        brew install redis
        if [ $? -ne 0 ]; then
            echo "Error: Failed to install Redis"
            exit 1
        fi
        brew services start redis
        if [ $? -ne 0 ]; then
            echo "Error: Failed to start Redis"
            exit 1
        fi
        echo "Redis installed successfully"
    fi
    echo "Redis service started successfully"
fi

# Check if .venv exists in suno_utils
if [ ! -d "./suno_utils/.venv" ]; then
    echo "Error: .venv directory not found in ./suno_utils"
    echo "Running uv sync..."
    (cd ./suno_utils && uv sync)
    if [ $? -ne 0 ]; then
        echo "Error: Failed to run uv sync"
        exit 1
    fi
    echo "uv sync completed successfully"
fi

# Check and confirm CALLBACK_HOST environment variable
echo -n "Current CALLBACK_HOST is: "
echo $CALLBACK_HOST
read -p "Is this CALLBACK_HOST value correct? (y/N) " callback_confirm
if [[ $callback_confirm != [yY] ]]; then
    echo "Please set the CALLBACK_HOST environment variable and try again"
    echo "Add to your ~/.zprofile: export CALLBACK_HOST=<your_machine_name>.han-mahi.ts.net"
    echo "Alternatively, you can set the CALLBACK_HOST environment variable in the studio_api/.env file"
    exit 1
fi

# Check if Modal config file exists
if [ ! -f "$HOME/.modal.toml" ]; then
    echo "Error: Modal config file not found at ~/.modal.toml"
    echo "Make sure you have access to Suno's Modal."
    (cd ./suno_utils && uv run modal token new)
    if [ $? -ne 0 ]; then
        echo "Error: Failed to create new Modal token."
        exit 1
    fi
    echo "Modal token created successfully."
fi

# Check if docker is running
if ! docker ps > /dev/null 2>&1; then
    echo "Error: Docker is not running"
    echo "Please start Docker Desktop and try again"
    exit 1
fi

if ! docker network ls | grep -q "elastic"; then
    docker network create elastic
fi

docker pull docker.elastic.co/elasticsearch/elasticsearch:8.15.2

# Check if elasticsearch container exists (running or stopped)
if docker ps -a --format '{{.Names}}' | grep -q "^elasticsearch$"; then
    # Container exists, check if it's running
    if ! docker ps --format '{{.Names}}' | grep -q "^elasticsearch$"; then
        echo "Starting elasticsearch container..."
        docker start elasticsearch
        if [ $? -ne 0 ]; then
            echo "Error: Failed to start elasticsearch container"
            exit 1
        fi
        echo "elasticsearch container started successfully"
    else
        echo "elasticsearch container is already running"
    fi
else
    # Container doesn't exist, create and run it
    echo "Creating elasticsearch container..."
    docker run -d --name elasticsearch --net elastic -p 9200:9200 -p 9300:9300 -e "discovery.type=single-node" -e "xpack.security.enabled=false" -t docker.elastic.co/elasticsearch/elasticsearch:8.15.2
    if [ $? -ne 0 ]; then
        echo "Error: Failed to create elasticsearch container"
        exit 1
    fi
    echo "elasticsearch container created and started successfully"
fi


echo "Setting up database and user..."
# Run psql commands to create database and user (idempotent)
psql -d template1 << EOF
-- Create databases if they don't exist
SELECT 'CREATE DATABASE suno_studio' WHERE NOT EXISTS (SELECT FROM pg_database WHERE datname = 'suno_studio')\gexec
SELECT 'CREATE DATABASE suno2' WHERE NOT EXISTS (SELECT FROM pg_database WHERE datname = 'suno2')\gexec

-- Create user if it doesn't exist
DO \$\$
BEGIN
  IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = 'suno') THEN
    CREATE USER suno WITH LOGIN PASSWORD 'suno';
  END IF;
END
\$\$;

-- Grant privileges (safe to run multiple times)
GRANT ALL PRIVILEGES ON DATABASE suno_studio TO suno;
GRANT ALL PRIVILEGES ON DATABASE suno2 TO suno;
ALTER USER suno CREATEDB;
EOF

# Grant schema privileges on each database
for db in suno_studio suno2; do
  psql -d $db << EOF
GRANT USAGE ON SCHEMA public TO suno;
GRANT CREATE ON SCHEMA public TO suno;
EOF
done

if [ $? -ne 0 ]; then
    echo "Error: Failed to set up database and user"
    exit 1
fi
echo "Database and user setup completed successfully"

echo "Migrating database..."
(cd ./studio_api && uv run manage.py migrate_all)
if [ $? -ne 0 ]; then
    echo "Error: Failed to migrate database"
    exit 1
fi
echo "Database migrated successfully"

echo "Seeding model data..."
(cd ./studio_api && uv run manage.py seed_model_data)
if [ $? -ne 0 ]; then
    echo "Error: Failed to seed model data"
    exit 1
fi
echo "Model data seeded successfully"

# Check if DynamoDB is already set up by looking for local DynamoDB container
if ! docker ps | grep -q "dynamodb-local"; then
    echo "Setting up DynamoDB..."
    # Set up local DynamoDB environment variables and run migration scripts
    export AWS_SECRET_ACCESS_KEY="local"
    export AWS_ACCESS_KEY_ID="local"

    # Run DynamoDB bootstrap script from studio_api directory
    echo "Running DynamoDB bootstrap script..."
    (cd ./studio_api && ./scripts/local_lyrics_migration/bootstrap.sh)
    if [ $? -ne 0 ]; then
        echo "Error: Failed to run DynamoDB bootstrap script"
        exit 1
    fi
    echo "DynamoDB bootstrap completed successfully"

    # Run DynamoDB migration script from studio_api directory
    echo "Running DynamoDB migration script..."
    (cd ./studio_api && uv run scripts/local_lyrics_migration/migrate_local_lyrics.py)
    if [ $? -ne 0 ]; then
        echo "Error: Failed to run DynamoDB migration script"
        exit 1
    fi
    echo "DynamoDB migration completed successfully"
else
    echo "DynamoDB is already running"
fi


# Prepare DSP engine artifacts
echo "Preparing DSP engine artifacts..."
(cd ./ui/app-ui && pnpm run prepare-dsp)
if [ $? -ne 0 ]; then
    echo "Error: Failed to prepare DSP engine artifacts"
    exit 1
fi
echo "DSP engine artifacts prepared successfully"


echo "Setup complete!"
echo "Run the following task to get started:"
echo "  - Sync Packages and Models"
echo "Then, go to Run/Debug and press: Run All"