#!/bin/bash
# Database setup for Trax project

set -e

# Color codes
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m'

echo -e "${BLUE}🗄️  Trax Database Setup${NC}"
echo "================================"

# Default database configuration
DB_NAME="${TRAX_DB_NAME:-trax}"
DB_USER="${TRAX_DB_USER:-postgres}"
DB_HOST="${TRAX_DB_HOST:-localhost}"
DB_PORT="${TRAX_DB_PORT:-5432}"
TEST_DB_NAME="${DB_NAME}_test"

# Parse arguments
ACTION=${1:-setup}

# Check PostgreSQL is installed
if ! command -v psql &> /dev/null; then
    echo -e "${RED}❌ PostgreSQL not installed${NC}"
    echo -e "${YELLOW}💡 Install with: brew install postgresql@15${NC}"
    exit 1
fi

# Check PostgreSQL is running
if ! pg_isready -h $DB_HOST -p $DB_PORT &> /dev/null; then
    echo -e "${RED}❌ PostgreSQL is not running${NC}"
    echo -e "${YELLOW}💡 Start with: brew services start postgresql@15${NC}"
    exit 1
fi

echo -e "${GREEN}✅ PostgreSQL is running${NC}"

case "$ACTION" in
    setup)
        echo -e "\n${CYAN}Creating databases...${NC}"
        
        # Create main database
        if psql -U $DB_USER -h $DB_HOST -lqt | cut -d \| -f 1 | grep -qw $DB_NAME; then
            echo -e "  ${YELLOW}⚠️  Database '$DB_NAME' already exists${NC}"
        else
            createdb -U $DB_USER -h $DB_HOST $DB_NAME
            echo -e "  ${GREEN}✅ Created database '$DB_NAME'${NC}"
        fi
        
        # Create test database
        if psql -U $DB_USER -h $DB_HOST -lqt | cut -d \| -f 1 | grep -qw $TEST_DB_NAME; then
            echo -e "  ${YELLOW}⚠️  Database '$TEST_DB_NAME' already exists${NC}"
        else
            createdb -U $DB_USER -h $DB_HOST $TEST_DB_NAME
            echo -e "  ${GREEN}✅ Created database '$TEST_DB_NAME'${NC}"
        fi
        
        # Create .env.local with database URL if it doesn't exist
        ENV_FILE="$(dirname "$0")/../.env.local"
        if [ ! -f "$ENV_FILE" ]; then
            echo -e "\n${CYAN}Creating .env.local...${NC}"
            cat > "$ENV_FILE" << EOF
# Trax Local Environment Variables
DATABASE_URL=postgresql://$DB_USER@$DB_HOST:$DB_PORT/$DB_NAME
TEST_DATABASE_URL=postgresql://$DB_USER@$DB_HOST:$DB_PORT/$TEST_DB_NAME

# Whisper Model Configuration
WHISPER_MODEL=distil-large-v3
WHISPER_DEVICE=cpu  # Change to 'mps' for M3 Mac optimization
WHISPER_COMPUTE_TYPE=int8_float32

# Batch Processing
BATCH_SIZE=10
MAX_WORKERS=4

# Cache Configuration
CACHE_TTL_SECONDS=3600
EOF
            echo -e "  ${GREEN}✅ Created .env.local with database configuration${NC}"
        else
            echo -e "  ${YELLOW}⚠️  .env.local already exists${NC}"
        fi
        
        # Initialize Alembic if not already done
        if [ ! -d "alembic" ]; then
            echo -e "\n${CYAN}Initializing Alembic...${NC}"
            cd "$(dirname "$0")/.."
            source .venv/bin/activate 2>/dev/null || true
            alembic init alembic
            echo -e "  ${GREEN}✅ Alembic initialized${NC}"
            
            # Update alembic.ini with database URL
            sed -i.bak "s|sqlalchemy.url = .*|sqlalchemy.url = postgresql://$DB_USER@$DB_HOST:$DB_PORT/$DB_NAME|" alembic.ini
            rm alembic.ini.bak
            echo -e "  ${GREEN}✅ Updated alembic.ini with database URL${NC}"
        else
            echo -e "  ${YELLOW}⚠️  Alembic already initialized${NC}"
        fi
        
        echo -e "\n${GREEN}✅ Database setup complete!${NC}"
        echo ""
        echo "Database URLs:"
        echo "  Main: postgresql://$DB_USER@$DB_HOST:$DB_PORT/$DB_NAME"
        echo "  Test: postgresql://$DB_USER@$DB_HOST:$DB_PORT/$TEST_DB_NAME"
        ;;
    
    drop)
        echo -e "\n${RED}⚠️  Dropping databases...${NC}"
        read -p "Are you sure? This will delete all data! (y/N): " confirm
        if [ "$confirm" = "y" ] || [ "$confirm" = "Y" ]; then
            dropdb -U $DB_USER -h $DB_HOST --if-exists $DB_NAME
            echo -e "  ${GREEN}✅ Dropped database '$DB_NAME'${NC}"
            dropdb -U $DB_USER -h $DB_HOST --if-exists $TEST_DB_NAME
            echo -e "  ${GREEN}✅ Dropped database '$TEST_DB_NAME'${NC}"
        else
            echo "  Cancelled"
        fi
        ;;
    
    reset)
        echo -e "\n${YELLOW}Resetting databases...${NC}"
        $0 drop
        $0 setup
        ;;
    
    status)
        echo -e "\n${CYAN}Database Status:${NC}"
        echo ""
        
        # Check main database
        if psql -U $DB_USER -h $DB_HOST -lqt | cut -d \| -f 1 | grep -qw $DB_NAME; then
            echo -e "  ${GREEN}✅ Main database '$DB_NAME' exists${NC}"
            
            # Show table count
            table_count=$(psql -U $DB_USER -h $DB_HOST -d $DB_NAME -t -c "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = 'public';" 2>/dev/null || echo "0")
            echo -e "    Tables: $table_count"
        else
            echo -e "  ${RED}❌ Main database '$DB_NAME' does not exist${NC}"
        fi
        
        # Check test database
        if psql -U $DB_USER -h $DB_HOST -lqt | cut -d \| -f 1 | grep -qw $TEST_DB_NAME; then
            echo -e "  ${GREEN}✅ Test database '$TEST_DB_NAME' exists${NC}"
        else
            echo -e "  ${RED}❌ Test database '$TEST_DB_NAME' does not exist${NC}"
        fi
        
        # Check Alembic
        if [ -d "alembic" ]; then
            echo -e "  ${GREEN}✅ Alembic is initialized${NC}"
            
            # Check current revision
            cd "$(dirname "$0")/.."
            source .venv/bin/activate 2>/dev/null || true
            current_rev=$(alembic current 2>/dev/null | grep -o '[a-f0-9]\{12\}' | head -1 || echo "none")
            echo -e "    Current revision: $current_rev"
        else
            echo -e "  ${YELLOW}⚠️  Alembic not initialized${NC}"
        fi
        ;;
    
    *)
        echo "Usage: $0 [setup|drop|reset|status]"
        echo ""
        echo "Commands:"
        echo "  setup  - Create databases and initialize Alembic"
        echo "  drop   - Drop all databases (requires confirmation)"
        echo "  reset  - Drop and recreate databases"
        echo "  status - Show database status"
        echo ""
        echo "Environment variables:"
        echo "  TRAX_DB_NAME - Database name (default: trax)"
        echo "  TRAX_DB_USER - Database user (default: postgres)"
        echo "  TRAX_DB_HOST - Database host (default: localhost)"
        echo "  TRAX_DB_PORT - Database port (default: 5432)"
        exit 1
        ;;
esac