#!/bin/bash

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

# Default values
ENVIRONMENT="local"
SERVICE=""
FOLLOW=false
TAIL_LINES=100
EXPORT_FILE=""
TIMESTAMPS=true

# Function to print usage
print_usage() {
    echo "Usage: $0 [OPTIONS]"
    echo ""
    echo "Options:"
    echo "  -e, --environment ENV    Environment (local|production) [default: local]"
    echo "  -s, --service SERVICE    Service name to filter logs"
    echo "  -f, --follow             Follow log output (like tail -f)"
    echo "  -n, --tail LINES         Number of lines to show from the end [default: 100]"
    echo "  -o, --output FILE        Export logs to file"
    echo "  -t, --no-timestamps      Don't show timestamps"
    echo "  -h, --help               Show this help message"
    echo ""
    echo "Examples:"
    echo "  $0                                    # Show last 100 lines of all services"
    echo "  $0 -f                                 # Follow all services logs"
    echo "  $0 -s backend-api                     # Show backend-api logs"
    echo "  $0 -s backend-api -f                  # Follow backend-api logs"
    echo "  $0 -e production -o logs.txt          # Export production logs to file"
    echo "  $0 -s postgres -n 50                  # Show last 50 lines of postgres"
}

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -e|--environment)
            ENVIRONMENT="$2"
            shift 2
            ;;
        -s|--service)
            SERVICE="$2"
            shift 2
            ;;
        -f|--follow)
            FOLLOW=true
            shift
            ;;
        -n|--tail)
            TAIL_LINES="$2"
            shift 2
            ;;
        -o|--output)
            EXPORT_FILE="$2"
            shift 2
            ;;
        -t|--no-timestamps)
            TIMESTAMPS=false
            shift
            ;;
        -h|--help)
            print_usage
            exit 0
            ;;
        *)
            echo -e "${RED}Unknown option: $1${NC}"
            print_usage
            exit 1
            ;;
    esac
done

# Determine compose file
if [ "$ENVIRONMENT" == "production" ]; then
    COMPOSE_FILE="docker-compose.production.yml"
else
    COMPOSE_FILE="docker-compose.yml"
fi

# Check if compose file exists
if [ ! -f "$COMPOSE_FILE" ]; then
    echo -e "${RED}❌ $COMPOSE_FILE not found. Aborting.${NC}" >&2
    exit 1
fi

# Build docker-compose logs command
COMPOSE_CMD="docker-compose -f $COMPOSE_FILE logs"

# Add timestamps option
if [ "$TIMESTAMPS" = true ]; then
    COMPOSE_CMD="$COMPOSE_CMD --timestamps"
fi

# Add tail option (only if not following)
if [ "$FOLLOW" = false ]; then
    COMPOSE_CMD="$COMPOSE_CMD --tail=$TAIL_LINES"
fi

# Add follow option
if [ "$FOLLOW" = true ]; then
    COMPOSE_CMD="$COMPOSE_CMD --follow"
fi

# Add service filter
if [ -n "$SERVICE" ]; then
    COMPOSE_CMD="$COMPOSE_CMD $SERVICE"
fi

# Function to display logs
display_logs() {
    if [ -n "$EXPORT_FILE" ]; then
        echo -e "${BLUE}📝 Exporting logs to $EXPORT_FILE...${NC}"
        eval "$COMPOSE_CMD" > "$EXPORT_FILE" 2>&1
        if [ $? -eq 0 ]; then
            echo -e "${GREEN}✅ Logs exported successfully to $EXPORT_FILE${NC}"
        else
            echo -e "${RED}❌ Failed to export logs${NC}"
            exit 1
        fi
    else
        # Display logs directly
        if [ -n "$SERVICE" ]; then
            echo -e "${BLUE}📋 Showing logs for service: $SERVICE${NC}"
        else
            echo -e "${BLUE}📋 Showing logs for all services${NC}"
        fi
        echo ""
        eval "$COMPOSE_CMD"
    fi
}

# Main execution
main() {
    # Display info
    echo -e "${BLUE}🔍 Collecting logs for ${ENVIRONMENT} environment${NC}"
    if [ -n "$SERVICE" ]; then
        echo -e "${BLUE}   Service: $SERVICE${NC}"
    fi
    if [ "$FOLLOW" = true ]; then
        echo -e "${BLUE}   Mode: Following logs (Ctrl+C to stop)${NC}"
    else
        echo -e "${BLUE}   Lines: Last $TAIL_LINES lines${NC}"
    fi
    echo ""
    
    # Display or export logs
    display_logs
}

# Run main function
main