#!/bin/bash

# Ringo log collection script
# Version: 1.1

set -e  # Exit on any error

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

# Function to print colored text
print_colored() {
    local color=$1
    local message=$2
    echo -e "${color}${message}${NC}"
}

# Check for sudo/root or Docker group membership
check_permissions() {
    if [[ $EUID -ne 0 ]] && ! groups | grep -q docker 2>/dev/null; then
        print_colored $RED "Error: Root privileges or Docker group membership required"
        print_colored $YELLOW "Hint: Run 'sudo -i' or add your user to the docker group"
        exit 1
    fi
}

# Check if Docker is installed
check_docker() {
    if ! command -v docker &> /dev/null; then
        print_colored $RED "Error: Docker is not installed or not in PATH"
        exit 1
    fi
}

# Display Docker containers
get_containers() {
    docker container ls --format "table {{.Names}}\t{{.Image}}\t{{.Status}}"
}

# Create log output directory
create_logs_dir() {
    local timestamp=$(date +"%Y%m%d_%H%M%S")
    LOGS_DIR="$HOME/ringo_logs_$timestamp"
    mkdir -p "$LOGS_DIR"
    print_colored $GREEN "Log directory created: $LOGS_DIR"
}

# Validate date format
validate_date() {
    local date_input=$1
    if [[ $date_input =~ ^[0-9]{4}-[0-9]{2}-[0-9]{2}$ ]]; then
        if date -d "$date_input" >/dev/null 2>&1; then
            return 0
        fi
    fi
    return 1
}

# Ask user for time range
get_time_range() {
    echo
    print_colored $BLUE "=== Time Range Selection ==="
    echo "1. Collect logs for the last 24 hours"
    echo "2. Collect logs for the last 7 days"
    echo "3. Specify custom date range"
    echo "4. Collect all available logs"
    echo
    read -p "Select an option [1-4]: " time_option
    
    case $time_option in
        1)
            SINCE_DATE=$(date -d "1 day ago" +"%Y-%m-%d")
            UNTIL_DATE=$(date +"%Y-%m-%d")
            ;;
        2)
            SINCE_DATE=$(date -d "7 days ago" +"%Y-%m-%d")
            UNTIL_DATE=$(date +"%Y-%m-%d")
            ;;
        3)
            while true; do
                read -p "Enter start date (YYYY-MM-DD): " SINCE_DATE
                if validate_date "$SINCE_DATE"; then
                    break
                else
                    print_colored $RED "Invalid date format. Use YYYY-MM-DD"
                fi
            done
            
            while true; do
                read -p "Enter end date (YYYY-MM-DD): " UNTIL_DATE
                if validate_date "$UNTIL_DATE"; then
                    break
                else
                    print_colored $RED "Invalid date format. Use YYYY-MM-DD"
                fi
            done
            ;;
        4)
            SINCE_DATE=""
            UNTIL_DATE=""
            ;;
        *)
            print_colored $RED "Invalid choice. Defaulting to last 24 hours"
            SINCE_DATE=$(date -d "1 day ago" +"%Y-%m-%d")
            UNTIL_DATE=$(date +"%Y-%m-%d")
            ;;
    esac
    
    if [[ -n "$SINCE_DATE" && -n "$UNTIL_DATE" ]]; then
        print_colored $GREEN "Selected range: $SINCE_DATE to $UNTIL_DATE"
    else
        print_colored $GREEN "Collecting all available logs"
    fi
}

# Collect logs from backend container
collect_backend_logs() {
    local container_name=$1
    print_colored $BLUE "Collecting logs from backend container: $container_name"
    
    if docker exec "$container_name" test -d /logs 2>/dev/null; then
        print_colored $GREEN "Found /logs folder in $container_name"
        print_colored $YELLOW "Available log files:"
        docker exec "$container_name" ls -la /logs/ || true
        
        if [[ -n "$SINCE_DATE" ]]; then
            docker exec "$container_name" find /logs -name "*.log*" -type f | while read -r logfile; do
                filename=$(basename "$logfile")
                if [[ $filename =~ [0-9]{4}-[0-9]{2}-[0-9]{2} ]]; then
                    file_date=$(echo "$filename" | grep -oE '[0-9]{4}-[0-9]{2}-[0-9]{2}')
                    if [[ "$file_date" > "$SINCE_DATE" || "$file_date" == "$SINCE_DATE" ]] && \
                       [[ "$file_date" < "$UNTIL_DATE" || "$file_date" == "$UNTIL_DATE" ]]; then
                        docker cp "$container_name:$logfile" "$LOGS_DIR/"
                        print_colored $GREEN "Copied: $filename"
                    fi
                else
                    docker cp "$container_name:$logfile" "$LOGS_DIR/"
                    print_colored $GREEN "Copied: $filename"
                fi
            done
        else
            docker cp "$container_name:/logs/." "$LOGS_DIR/"
            print_colored $GREEN "All logs copied from $container_name:/logs/"
        fi
    else
        print_colored $YELLOW "/logs not found in container $container_name, falling back to docker logs"
        collect_container_logs "$container_name"
    fi
}

# Collect logs via `docker logs`
collect_container_logs() {
    local container_name=$1
    local log_file="$LOGS_DIR/${container_name}.log"
    
    print_colored $BLUE "Collecting logs from container: $container_name"
    
    if [[ -n "$SINCE_DATE" && -n "$UNTIL_DATE" ]]; then
        docker logs "$container_name" -t --since="$SINCE_DATE" --until="$UNTIL_DATE" &> "$log_file"
    else
        docker logs "$container_name" -t &> "$log_file"
    fi
    
    if [[ -s "$log_file" ]]; then
        print_colored $GREEN "Logs saved to: $log_file"
    else
        print_colored $YELLOW "Log is empty or inaccessible for $container_name"
        rm -f "$log_file"
    fi
}

# Identify Ringo containers
is_ringo_container() {
    local container_name=$1
    if [[ $container_name =~ backend ]] || \
       [[ $container_name =~ micromdm ]] || \
       [[ $container_name =~ roll-deploy-db ]] || \
       [[ $container_name =~ ringo-letsencrypt ]] || \
       [[ $container_name =~ ringo-nginx-proxy ]]; then
        return 0
    fi
    return 1
}

# Select containers for log collection
select_containers() {
    echo
    print_colored $BLUE "=== Available Docker Containers ==="
    
    local containers=($(docker container ls --format "{{.Names}}"))
    
    if [[ ${#containers[@]} -eq 0 ]]; then
        print_colored $RED "No running containers found"
        exit 1
    fi
    
    echo "Running containers:"
    get_containers
    echo
    
    local backend_containers=()
    local micromdm_containers=()
    local database_containers=()
    local proxy_containers=()
    local letsencrypt_containers=()
    local other_containers=()
    
    for container in "${containers[@]}"; do
        if [[ $container =~ backend ]]; then
            backend_containers+=("$container")
        elif [[ $container =~ micromdm ]]; then
            micromdm_containers+=("$container")
        elif [[ $container =~ roll-deploy-db ]]; then
            database_containers+=("$container")
        elif [[ $container =~ ringo-nginx-proxy ]]; then
            proxy_containers+=("$container")
        elif [[ $container =~ ringo-letsencrypt ]]; then
            letsencrypt_containers+=("$container")
        elif is_ringo_container "$container"; then
            other_containers+=("$container")
        else
            other_containers+=("$container")
        fi
    done
    
    echo "Detected Ringo containers:"
    [[ ${#backend_containers[@]} -gt 0 ]] && print_colored $GREEN "Backend: ${backend_containers[*]}"
    [[ ${#micromdm_containers[@]} -gt 0 ]] && print_colored $GREEN "MicroMDM: ${micromdm_containers[*]}"
    [[ ${#database_containers[@]} -gt 0 ]] && print_colored $GREEN "Database: ${database_containers[*]}"
    [[ ${#proxy_containers[@]} -gt 0 ]] && print_colored $GREEN "Nginx Proxy: ${proxy_containers[*]}"
    [[ ${#letsencrypt_containers[@]} -gt 0 ]] && print_colored $GREEN "Let's Encrypt: ${letsencrypt_containers[*]}"
    
    local ringo_containers=("${backend_containers[@]}" "${micromdm_containers[@]}" "${database_containers[@]}" "${proxy_containers[@]}" "${letsencrypt_containers[@]}")
    
    local non_ringo_containers=()
    for container in "${containers[@]}"; do
        if ! is_ringo_container "$container"; then
            non_ringo_containers+=("$container")
        fi
    done
    [[ ${#non_ringo_containers[@]} -gt 0 ]] && print_colored $YELLOW "Other containers: ${non_ringo_containers[*]}"
    
    echo
    echo "Choose log collection mode:"
    echo "1. Collect logs from all Ringo containers (recommended)"
    echo "2. Select containers manually"
    echo "3. Collect logs from all running containers"
    echo
    read -p "Select an option [1-3]: " collection_mode
    
    case $collection_mode in
        1)
            SELECTED_CONTAINERS=("${ringo_containers[@]}")
            ;;
        2)
            echo "Available containers:"
            for i in "${!containers[@]}"; do
                local marker=""
                if is_ringo_container "${containers[i]}"; then
                    marker=" (Ringo)"
                fi
                echo "$((i+1)). ${containers[i]}$marker"
            done
            echo
            read -p "Enter container numbers (e.g., 1 3 5): " selection
            SELECTED_CONTAINERS=()
            for num in $selection; do
                if [[ $num -ge 1 && $num -le ${#containers[@]} ]]; then
                    SELECTED_CONTAINERS+=("${containers[$((num-1))]}")
                fi
            done
            ;;
        3)
            SELECTED_CONTAINERS=("${containers[@]}")
            ;;
        *)
            print_colored $RED "Invalid selection. Defaulting to automatic mode"
            SELECTED_CONTAINERS=("${ringo_containers[@]}")
            ;;
    esac
    
    if [[ ${#SELECTED_CONTAINERS[@]} -eq 0 ]]; then
        print_colored $RED "No containers selected"
        exit 1
    fi
    
    print_colored $GREEN "Selected containers: ${SELECTED_CONTAINERS[*]}"
}

# Create archive from collected logs
create_archive() {
    echo
    print_colored $BLUE "=== Creating archive ==="
    
    local archive_name="ringo_logs_$(date +"%Y%m%d_%H%M%S").tar.gz"
    local archive_path="$HOME/$archive_name"
    
    cd "$(dirname "$LOGS_DIR")"
    tar -czf "$archive_path" "$(basename "$LOGS_DIR")"
    
    if [[ -f "$archive_path" ]]; then
        local archive_size=$(du -h "$archive_path" | cut -f1)
        print_colored $GREEN "Archive created: $archive_path (size: $archive_size)"
        
        echo
        print_colored $YELLOW "To send to Ringo support, use:"
        print_colored $BLUE "Email: support@ringomdm.ru"
        print_colored $BLUE "Archive: $archive_path"
    else
        print_colored $RED "Failed to create archive"
    fi
}

# Show collected log summary
show_summary() {
    echo
    print_colored $BLUE "=== Log Summary ==="
    
    if [[ -d "$LOGS_DIR" ]]; then
        print_colored $GREEN "Log directory: $LOGS_DIR"
        echo "Collected files:"
        du -h "$LOGS_DIR"/* 2>/dev/null | while read -r size file; do
            print_colored $YELLOW "  $size - $(basename "$file")"
        done
        
        local total_size=$(du -sh "$LOGS_DIR" | cut -f1)
        print_colored $GREEN "Total size: $total_size"
    fi
}

# Main function
main() {
    print_colored $BLUE "=== Ringo Log Collection Script ==="
    echo
    
    check_permissions
    check_docker
    create_logs_dir
    select_containers
    get_time_range
    
    echo
    print_colored $BLUE "=== Starting log collection ==="
    
    for container in "${SELECTED_CONTAINERS[@]}"; do
        if [[ $container =~ backend ]]; then
            collect_backend_logs "$container"
        else
            collect_container_logs "$container"
        fi
        echo
    done
    
    show_summary
    
    echo
    read -p "Create archive from logs? [y/N]: " create_arch
    if [[ $create_arch =~ ^[Yy] ]]; then
        create_archive
    fi
    
    print_colored $GREEN "Log collection completed successfully!"
}

# Run main
main "$@"