#!/bin/bash

# GremlinModes Agent Selection Script
# This script allows users to select which agent to use
# with the option to add solo agents to teams

# Text formatting
BOLD='\033[1m'
GREMLIN='\033[1;38;2;31;204;0m'
GREEN='\033[0;32m'
BLUE='\033[0;34m'
RED='\033[0;31m'
YELLOW='\033[0;33m'
NC='\033[0m' # No Color

# Banner
echo -e "${BOLD}${GREMLIN}"
echo "   ▄▄ • ▄▄▄  ▄▄▄ .• ▌ ▄ ·. ▄▄▌  ▪   ▐ ▄ • ▌ ▄ ·.       ·▄▄▄▄  ▄▄▄ ..▄▄ · "
echo "  ▐█ ▀ ▪▀▄ █·▀▄.▀··██ ▐███▪██•  ██ •█▌▐█·██ ▐███▪▪     ██▪ ██ ▀▄.▀·▐█ ▀. "
echo "  ▄█ ▀█▄▐▀▀▄ ▐▀▀▪▄▐█ ▌▐▌▐█·██▪  ▐█·▐█▐▐▌▐█ ▌▐▌▐█· ▄█▀▄ ▐█· ▐█▌▐▀▀▪▄▄▀▀▀█▄"
echo "  ▐█▄▪▐█▐█•█▌▐█▄▄▌██ ██▌▐█▌▐█▌▐▌▐█▌██▐█▌██ ██▌▐█▌▐█▌.▐▌██. ██ ▐█▄▄▌▐█▄▪▐█"
echo "  ·▀▀▀▀ .▀  ▀ ▀▀▀ ▀▀  █▪▀▀▀.▀▀▀ ▀▀▀▀▀ █▪▀▀  █▪▀▀▀ ▀█▄▀▪▀▀▀▀▀•  ▀▀▀  ▀▀▀▀ "
echo -e "${NC}"
echo -e "${BOLD}Agent Selection${NC}"
echo ""

# Check if jq is installed
if ! command -v jq &> /dev/null; then
    echo -e "${YELLOW}Warning: 'jq' is not installed. It's recommended for JSON processing.${NC}"
    echo -e "${YELLOW}You can install it using your package manager (e.g., apt install jq, brew install jq)${NC}"
    echo -e "${YELLOW}Continuing with basic functionality...${NC}"
    echo ""
    HAS_JQ=false
else
    HAS_JQ=true
fi

# Check if .gremlinmodes directory exists
if [ ! -d ".gremlinmodes" ]; then
    echo -e "${RED}Error: .gremlinmodes directory not found. Please run install.sh first.${NC}"
    exit 1
fi

# Check if agents directory exists
if [ ! -d ".gremlinmodes/agents" ]; then
    echo -e "${RED}Error: No agent configurations found. Please run install.sh first.${NC}"
    exit 1
fi

# Count available agents
agent_count=$(ls -1 .gremlinmodes/agents/*.json 2>/dev/null | wc -l)

if [ "$agent_count" -eq 0 ]; then
    echo -e "${RED}Error: No agent configuration files (.json) found in .gremlinmodes/agents/${NC}"
    exit 1
fi

# Function to check if a JSON file contains multiple agents (is a team)
is_team_file() {
    local file=$1
    if [ "$HAS_JQ" = true ]; then
        # Count the number of customModes in the JSON
        local mode_count=$(jq '.customModes | length' "$file" 2>/dev/null)
        
        # Ensure mode_count is a number, default to 0 if empty or invalid
        if [ -z "$mode_count" ] || ! [[ "$mode_count" =~ ^[0-9]+$ ]]; then
            mode_count=0
        fi
        
        if [ "$mode_count" -gt 1 ]; then
            return 0  # true - is a team
        else
            return 1  # false - is solo
        fi
    else
        # Fallback: use grep to check for multiple agent names (less reliable but works without jq)
        local name_count=$(grep -c '"name"' "$file" 2>/dev/null || echo "0")
        if [ "$name_count" -gt 1 ]; then
            return 0  # true - is a team
        else
            return 1  # false - is solo
        fi
    fi
}

# Store file paths
declare -a files
declare -a team_files
declare -a solo_files

# Categorize files
declare -a dev_teams
declare -a specialized_teams
declare -a solo_agents
declare -a uncategorized_agents

# Analyze each file
for file in .gremlinmodes/agents/*.json; do
    if [ -f "$file" ]; then
        filename=$(basename "$file")
        
        # Check file patterns and content
        if [[ "$filename" == *-solo-* ]]; then
            # Explicitly marked as solo
            solo_agents+=("$file")
            solo_files+=("$file")
        elif [[ "$filename" == *-dev-* ]]; then
            # Development team pattern
            dev_teams+=("$file")
            team_files+=("$file")
        elif [[ "$filename" == *-team-* ]] || [[ "$filename" == *-migration-* ]]; then
            # Team pattern or migration pattern
            specialized_teams+=("$file")
            team_files+=("$file")
        elif is_team_file "$file"; then
            # Check content: if multiple agents, it's a team
            specialized_teams+=("$file")
            team_files+=("$file")
        elif [[ "$filename" == *-seo-* ]] || [[ "$filename" == *-product-* ]] || \
             [[ "$filename" == *-nextjs-* ]] || [[ "$filename" == *-ios-* ]] || \
             [[ "$filename" == *-docs-* ]] || [[ "$filename" == *-design-* ]] || \
             [[ "$filename" == *-audit-* ]] || [[ "$filename" == *-testing-* ]] || \
             [[ "$filename" == *-rebuild-* ]]; then
            # Known specialized team patterns
            specialized_teams+=("$file")
            team_files+=("$file")
        else
            # Uncategorized - could be solo or team
            uncategorized_agents+=("$file")
        fi
    fi
done

# Display agents by category
echo -e "${BLUE}Available Agents:${NC}"
echo ""

count=0
team_count=0
solo_count=0

# Development Teams
echo -e "${YELLOW}Development Teams:${NC}"
if [ ${#dev_teams[@]} -eq 0 ]; then
    echo "  None available"
else
    for file in "${dev_teams[@]}"; do
        filename=$(basename "$file")
        count=$((count + 1))
        team_count=$((team_count + 1))
        echo "  ${count}. ${filename%.json}"
        files[$((count-1))]="$file"
    done
fi
echo ""

# Specialized Teams
echo -e "${YELLOW}Specialized Teams:${NC}"
if [ ${#specialized_teams[@]} -eq 0 ]; then
    echo "  None available"
else
    for file in "${specialized_teams[@]}"; do
        filename=$(basename "$file")
        count=$((count + 1))
        team_count=$((team_count + 1))
        echo "  ${count}. ${filename%.json}"
        files[$((count-1))]="$file"
    done
fi
echo ""

# Solo Agents
echo -e "${YELLOW}Solo Gremlins:${NC}"
if [ ${#solo_agents[@]} -eq 0 ]; then
    echo "  None available"
else
    for file in "${solo_agents[@]}"; do
        filename=$(basename "$file")
        count=$((count + 1))
        solo_count=$((solo_count + 1))
        echo "  ${count}. ${filename%.json}"
        files[$((count-1))]="$file"
    done
fi
echo ""

# Uncategorized Agents (if any)
if [ ${#uncategorized_agents[@]} -gt 0 ]; then
    echo -e "${YELLOW}Other Agents:${NC}"
    for file in "${uncategorized_agents[@]}"; do
        filename=$(basename "$file")
        count=$((count + 1))
        echo "  ${count}. ${filename%.json}"
        files[$((count-1))]="$file"
        
        # Add to appropriate array based on content check
        if is_team_file "$file"; then
            team_files+=("$file")
            team_count=$((team_count + 1))
        else
            solo_files+=("$file")
            solo_count=$((solo_count + 1))
        fi
    done
    echo ""
fi

# Count all agents (before adding the special option)
total_count=${#files[@]}

# Add the "All Gremlins" option
all_option_number=$((total_count + 1))
echo -e "${YELLOW}Special Options:${NC}"
echo "  ${all_option_number}. All Gremlins (All Teams & All Solos)"
echo ""

# Prompt user to select an agent
echo -e "${BLUE}Select a base agent, team, or special option by number (1-$all_option_number):${NC}"
read -p "> " selection

# Validate selection
if ! [[ "$selection" =~ ^[0-9]+$ ]]; then
    echo -e "${RED}Error: Please enter a number.${NC}"
    exit 1
fi

if [ "$selection" -lt 1 ] || [ "$selection" -gt "$all_option_number" ]; then
    echo -e "${RED}Error: Selection out of range.${NC}"
    exit 1
fi

# Check if the special "ALL" option was selected
if [ "$selection" -eq "$all_option_number" ]; then
    echo -e "${GREEN}Selected: ${BOLD}All Gremlins (All Teams & All Solos)${NC} - The Ultimate Gremlin Army!${NC}"
    
    if [ "$HAS_JQ" = false ]; then
        echo -e "${RED}Error: 'jq' is required for merging all agents. Please install jq first.${NC}"
        exit 1
    fi
    
    # Create a temporary file with an empty customModes array
    temp_file=$(mktemp)
    echo '{"customModes":[]}' > "$temp_file"
    
    # First, add all teams
    echo -e "${BLUE}Assembling all teams...${NC}"
    for team_file in "${team_files[@]}"; do
        team_name=$(basename "$team_file" .json)
        jq -s '.[0].customModes = (.[0].customModes + .[1].customModes) | .[0]' "$temp_file" "$team_file" > "$temp_file.new"
        mv "$temp_file.new" "$temp_file"
        echo -e "${GREEN}  ✓ Added team: ${BOLD}$team_name${NC}"
    done
    
    # Then, add all solo agents
    echo -e "${BLUE}Recruiting all solo agents...${NC}"
    for solo_file in "${solo_files[@]}"; do
        solo_name=$(basename "$solo_file" .json)
        jq -s '.[0].customModes = (.[0].customModes + .[1].customModes) | .[0]' "$temp_file" "$solo_file" > "$temp_file.new"
        mv "$temp_file.new" "$temp_file"
        echo -e "${GREEN}  ✓ Added solo: ${BOLD}$solo_name${NC}"
    done
    
    echo -e "${GREEN}Successfully assembled the Mega Gremlin Squad with ${team_count} teams and ${solo_count} solo agents!${NC}"
    
else
    # Regular selection (existing logic)
    # Get the selected agent file
    selected_file="${files[$((selection-1))]}"
    selected_name=$(basename "$selected_file" .json)

    # Check if the selected agent is a team
    is_team=false
    for team_file in "${team_files[@]}"; do
        if [ "$selected_file" = "$team_file" ]; then
            is_team=true
            break
        fi
    done

    # If a solo agent was selected, just use it directly
    if [ "$is_team" = false ]; then
        echo -e "${GREEN}Selected solo agent: ${BOLD}$selected_name${NC}"

        # Create a temporary file for the final configuration
        temp_file=$(mktemp)
        cp "$selected_file" "$temp_file"

    else
        # Offer to add solo agents to the team
        echo -e "${GREEN}Selected team: ${BOLD}$selected_name${NC}"

        if [ $solo_count -eq 0 ]; then
            echo -e "${YELLOW}No solo agents available to add to the team.${NC}"

            # Create a temporary file for the final configuration
            temp_file=$(mktemp)
            cp "$selected_file" "$temp_file"

        else
            echo -e "${BLUE}Would you like to add solo agents to this team? (y/n)${NC}"
            read -p "> " add_solos

            if [[ "$add_solos" =~ ^[Yy]$ ]]; then
                # Create a temporary file for the merged configuration
                temp_file=$(mktemp)
                cp "$selected_file" "$temp_file"

                # Display available solo agents
                echo -e "${BLUE}Available Solo Agents:${NC}"
                for ((i=0; i<$solo_count; i++)); do
                    filename=$(basename "${solo_files[$i]}" .json)
                    echo "  $((i+1)). $filename"
                done

                # Allow user to select multiple solo agents
                echo -e "${BLUE}Select solo agents to add:${NC}"
                echo -e "${BLUE}  1-${solo_count}: Individual solo agents by number${NC}"
                echo -e "${BLUE}  all: Add all solo agents${NC}"
                echo -e "${BLUE}  0: None (cancel)${NC}"
                echo -e "${BLUE}Enter option or numbers (space-separated for multiple):${NC}"
                read -p "> " solo_selections

                if [ "$solo_selections" = "0" ]; then
                    echo -e "${YELLOW}No solo agents selected.${NC}"
                elif [ "$solo_selections" = "all" ]; then
                    # Add all solo agents
                    echo -e "${GREEN}Adding all solo agents to the team${NC}"

                    for ((i=0; i<$solo_count; i++)); do
                        solo_file="${solo_files[$i]}"
                        solo_name=$(basename "$solo_file" .json)

                        if [ "$HAS_JQ" = true ]; then
                            # Use jq to merge the JSON files
                            jq -s '.[0].customModes = (.[0].customModes + .[1].customModes) | .[0]' "$temp_file" "$solo_file" > "$temp_file.new"
                            mv "$temp_file.new" "$temp_file"
                            echo -e "${GREEN}Added ${BOLD}$solo_name${NC}${GREEN} to the team${NC}"
                        else
                            echo -e "${YELLOW}Cannot add $solo_name - jq is required for JSON merging${NC}"
                        fi
                    done
                else
                    # Add selected solo agents
                    for num in $solo_selections; do
                        if ! [[ "$num" =~ ^[0-9]+$ ]]; then
                            echo -e "${RED}Warning: '$num' is not a number, skipping${NC}"
                            continue
                        fi

                        if [ "$num" -lt 1 ] || [ "$num" -gt "$solo_count" ]; then
                            echo -e "${RED}Warning: '$num' is out of range, skipping${NC}"
                            continue
                        fi

                        solo_file="${solo_files[$((num-1))]}"
                        solo_name=$(basename "$solo_file" .json)

                        if [ "$HAS_JQ" = true ]; then
                            # Use jq to merge the JSON files
                            jq -s '.[0].customModes = (.[0].customModes + .[1].customModes) | .[0]' "$temp_file" "$solo_file" > "$temp_file.new"
                            mv "$temp_file.new" "$temp_file"
                            echo -e "${GREEN}Added ${BOLD}$solo_name${NC}${GREEN} to the team${NC}"
                        else
                            echo -e "${YELLOW}Cannot add $solo_name - jq is required for JSON merging${NC}"
                        fi
                    done
                fi
            else
                # No solo agents to add, just use the selected team
                echo -e "${GREEN}Using team without additional solo agents${NC}"
                temp_file=$(mktemp)
                cp "$selected_file" "$temp_file"
            fi
        fi
    fi
fi

# Backup existing .roomodes file if it exists
if [ -f ".roomodes" ]; then
    # Create timestamp for backup directory
    timestamp=$(date +%Y%m%d%H%M%S)

    # Create backup directory structure if it doesn't exist
    backup_dir=".gremlinmodes/backups/roomodes"
    mkdir -p "$backup_dir"

    # Backup filename with timestamp
    backup_file="$backup_dir/.roomodes.$timestamp"

    # Copy the file (using cp instead of mv to ensure original permissions remain intact)
    cp ".roomodes" "$backup_file"

    echo -e "${BLUE}Backed up existing .roomodes to $backup_file${NC}"
fi

# Copy the final configuration to .roomodes
cp "$temp_file" ".roomodes"

# Clean up the temporary file
rm -f "$temp_file"

echo -e "${GREEN}Agent configuration activated and copied to .roomodes${NC}"
echo -e "${BLUE}You can now use this agent with Roo Code in your editor.${NC}"

echo ""
echo -e "${GREMLIN}Gremlin team is ready to cause strategic chaos! 😈${NC}"