#!/bin/bash
# Cleanup script for fetch_ml deployments
# Removes orphaned containers, volumes, and networks

set -e

REPO_ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)

# Colors
GREEN='\033[0;32m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
NC='\033[0m'

print_success() { echo -e "${GREEN}[OK]${NC} $1"; }
print_error() { echo -e "${RED}[ERROR]${NC} $1"; }
print_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
print_info() { echo "[INFO] $1"; }

show_usage() {
    echo "Usage: $0 [ENVIRONMENT|all]"
    echo ""
    echo "Environments:"
    echo "  dev         Clean up development environment"
    echo "  staging     Clean up staging environment"
    echo "  prod        Clean up production environment"
    echo "  homelab     Clean up homelab environment"
    echo "  all         Clean up all environments (default)"
    echo ""
    echo "Examples:"
    echo "  $0 dev       # Clean dev only"
    echo "  $0 all       # Clean everything"
}

cleanup_environment() {
    local env=$1
    local compose_file

    case $env in
        dev)
            compose_file="${REPO_ROOT}/deployments/docker-compose.dev.yml"
            ;;
        staging)
            compose_file="${REPO_ROOT}/deployments/docker-compose.staging.yml"
            ;;
        prod)
            compose_file="${REPO_ROOT}/deployments/docker-compose.prod.yml"
            ;;
        homelab)
            compose_file="${REPO_ROOT}/deployments/docker-compose.homelab-secure.yml"
            ;;
        *)
            print_error "Unknown environment: $env"
            show_usage
            exit 1
            ;;
    esac

    if [ ! -f "$compose_file" ]; then
        print_warn "Compose file not found: $compose_file"
        return 0
    fi

    print_info "Cleaning up $env environment..."

    # Stop and remove containers, networks, volumes
    docker-compose -f "$compose_file" down -v --remove-orphans 2>/dev/null || true

    # Remove any leftover containers with project name
    local project_name
    project_name=$(basename "$compose_file" .yml | sed 's/docker-compose\.//')
    local containers
    containers=$(docker ps -aq --filter "name=ml-${project_name}" --filter "name=ml-experiments" 2>/dev/null || true)

    if [ -n "$containers" ]; then
        print_info "Removing leftover containers..."
        docker rm -f $containers 2>/dev/null || true
    fi

    print_success "${env} environment cleaned"
}

cleanup_all() {
    print_info "Cleaning up all FetchML environments..."

    # Clean each environment
    for env in dev staging prod homelab; do
        cleanup_environment "$env" || true
    done

    # Remove any remaining fetchml containers by name pattern
    local all_containers
    all_containers=$(docker ps -aq --filter "name=^ml-" 2>/dev/null || true)

    if [ -n "$all_containers" ]; then
        print_info "Removing remaining FetchML containers..."
        docker rm -f $all_containers 2>/dev/null || true
    fi

    # Only prune FetchML-specific networks (not all unused)
    local networks
    networks=$(docker network ls -q --filter "name=^ml-" 2>/dev/null || true)
    if [ -n "$networks" ]; then
        print_info "Removing FetchML networks..."
        docker network rm $networks 2>/dev/null || true
    fi

    print_success "All FetchML environments cleaned (other Docker data preserved)"
}

main() {
    local env=${1:-all}

    case $env in
        dev|staging|prod|homelab)
            cleanup_environment "$env"
            ;;
        all)
            cleanup_all
            ;;
        -h|--help)
            show_usage
            exit 0
            ;;
        *)
            print_error "Unknown option: $env"
            show_usage
            exit 1
            ;;
    esac
}

main "$@"