#!/bin/bash
# MEV Bot Development Environment Manager
# Supports branch selection and podman-in-podman development

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"

# Default values
DEFAULT_BRANCH="master-dev"
COMPOSE_FILE="docker-compose.dev.yml"

# Function to print colored messages
print_info() {
    echo -e "${BLUE}ℹ${NC} $1"
}

print_success() {
    echo -e "${GREEN}✅${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}⚠${NC} $1"
}

print_error() {
    echo -e "${RED}❌${NC} $1"
}

# Function to show usage
usage() {
    cat << EOF
${BLUE}MEV Bot Development Environment Manager${NC}

Usage: $0 [COMMAND] [OPTIONS]

Commands:
    start [BRANCH]      Start development environment with specified branch (default: master-dev)
    stop                Stop development environment
    restart [BRANCH]    Restart development environment with specified branch
    rebuild [BRANCH]    Rebuild and restart with specified branch
    logs [OPTIONS]      Show logs (pass options like -f for follow)
    shell               Open shell in running container
    status              Show container status
    clean               Clean all containers and volumes
    branches            List available git branches
    switch BRANCH       Switch to a different branch and rebuild

Options:
    -h, --help          Show this help message

Examples:
    $0 start master-dev                 # Start with master-dev branch
    $0 start fix-critical-arbitrage-bugs # Start with feature branch
    $0 rebuild master                    # Rebuild with master branch
    $0 logs -f                          # Follow logs
    $0 switch feat-new-feature          # Switch to new feature branch

Environment Variables:
    GIT_BRANCH          Git branch to use (default: master-dev)
    LOG_LEVEL           Logging level (default: debug)
    METRICS_ENABLED     Enable metrics (default: true)

EOF
}

# Function to check dependencies
check_dependencies() {
    if ! command -v podman-compose &> /dev/null; then
        print_error "podman-compose is not installed"
        print_info "Install with: pip3 install podman-compose"
        exit 1
    fi
}

# Function to list available branches
list_branches() {
    print_info "Available branches:"
    cd "$PROJECT_ROOT"
    git branch -a | grep -v HEAD | sed 's/^..../  /'
}

# Function to start development environment
start_dev() {
    local branch="${1:-$DEFAULT_BRANCH}"

    print_info "Starting MEV Bot development environment on branch: $branch"

    cd "$PROJECT_ROOT"

    # Export branch for docker-compose
    export GIT_BRANCH="$branch"

    # Build and start
    print_info "Building image for branch: $branch"
    podman-compose -f "$COMPOSE_FILE" build --build-arg GIT_BRANCH="$branch"

    print_info "Starting container..."
    podman-compose -f "$COMPOSE_FILE" up -d

    print_success "Development environment started on branch: $branch"
    print_info "View logs with: $0 logs -f"
    print_info "Access shell with: $0 shell"
}

# Function to stop development environment
stop_dev() {
    print_info "Stopping development environment..."
    cd "$PROJECT_ROOT"
    podman-compose -f "$COMPOSE_FILE" down
    print_success "Development environment stopped"
}

# Function to restart development environment
restart_dev() {
    local branch="${1:-$DEFAULT_BRANCH}"
    stop_dev
    start_dev "$branch"
}

# Function to rebuild development environment
rebuild_dev() {
    local branch="${1:-$DEFAULT_BRANCH}"

    print_info "Rebuilding development environment on branch: $branch"

    cd "$PROJECT_ROOT"
    export GIT_BRANCH="$branch"

    # Stop existing containers
    podman-compose -f "$COMPOSE_FILE" down

    # Rebuild with no cache
    print_info "Building image (no cache) for branch: $branch"
    podman-compose -f "$COMPOSE_FILE" build --no-cache --build-arg GIT_BRANCH="$branch"

    # Start
    print_info "Starting container..."
    podman-compose -f "$COMPOSE_FILE" up -d

    print_success "Development environment rebuilt and started on branch: $branch"
}

# Function to show logs
show_logs() {
    cd "$PROJECT_ROOT"
    podman-compose -f "$COMPOSE_FILE" logs "$@"
}

# Function to open shell
open_shell() {
    cd "$PROJECT_ROOT"
    local container=$(podman-compose -f "$COMPOSE_FILE" ps -q | head -1)
    if [ -z "$container" ]; then
        print_error "No running container found"
        exit 1
    fi
    print_info "Opening shell in container..."
    podman exec -it "$container" /bin/bash || podman exec -it "$container" /bin/sh
}

# Function to show status
show_status() {
    cd "$PROJECT_ROOT"
    print_info "Development environment status:"
    podman-compose -f "$COMPOSE_FILE" ps
}

# Function to clean everything
clean_all() {
    print_warning "This will remove all containers, images, and volumes. Continue? (y/N)"
    read -r response
    if [[ "$response" =~ ^[Yy]$ ]]; then
        cd "$PROJECT_ROOT"
        podman-compose -f "$COMPOSE_FILE" down -v
        podman images | grep mev-bot | awk '{print $3}' | xargs -r podman rmi -f
        print_success "Cleaned all containers, images, and volumes"
    else
        print_info "Cancelled"
    fi
}

# Function to switch branches
switch_branch() {
    local branch="$1"

    if [ -z "$branch" ]; then
        print_error "Please specify a branch"
        usage
        exit 1
    fi

    print_info "Switching to branch: $branch"

    cd "$PROJECT_ROOT"

    # Check if branch exists
    if ! git rev-parse --verify "$branch" &> /dev/null; then
        print_error "Branch '$branch' does not exist"
        list_branches
        exit 1
    fi

    # Stop current environment
    stop_dev

    # Checkout branch
    print_info "Checking out branch: $branch"
    git checkout "$branch"

    # Rebuild and start
    rebuild_dev "$branch"
}

# Main script logic
main() {
    check_dependencies

    case "${1:-}" in
        start)
            start_dev "${2:-$DEFAULT_BRANCH}"
            ;;
        stop)
            stop_dev
            ;;
        restart)
            restart_dev "${2:-$DEFAULT_BRANCH}"
            ;;
        rebuild)
            rebuild_dev "${2:-$DEFAULT_BRANCH}"
            ;;
        logs)
            shift
            show_logs "$@"
            ;;
        shell)
            open_shell
            ;;
        status)
            show_status
            ;;
        clean)
            clean_all
            ;;
        branches)
            list_branches
            ;;
        switch)
            switch_branch "$2"
            ;;
        -h|--help|help)
            usage
            ;;
        *)
            print_error "Unknown command: ${1:-}"
            echo ""
            usage
            exit 1
            ;;
    esac
}

# Run main function
main "$@"