#!/bin/bash
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

echo -e "${GREEN}=== k3s Insecure Registry Configuration Script ===${NC}"
echo ""
echo "This script will configure your k3s nodes to allow pulling images"
echo "from your Gitea registry at 192.168.1.208:3002"
echo ""

# Get the directory where this script is located
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"

# Check if registries.yaml.example exists
if [ ! -f "$SCRIPT_DIR/registries.yaml.example" ]; then
    echo -e "${RED}Error: registries.yaml.example not found!${NC}"
    exit 1
fi

echo -e "${YELLOW}Step 1: Configure k3s Nodes${NC}"
echo ""
echo "You need to configure the following on EACH k3s node:"
echo "  1. Copy registries.yaml to /etc/rancher/k3s/registries.yaml"
echo "  2. Restart k3s or k3s-agent service"
echo ""

# Ask user for node IPs
echo -e "${YELLOW}Enter your k3s node IP addresses:${NC}"
echo "(Press Enter after each IP, then type 'done' when finished)"
echo ""

NODE_IPS=()
while true; do
    read -p "Node IP (or 'done'): " node_ip
    if [ "$node_ip" = "done" ]; then
        break
    fi
    if [ -n "$node_ip" ]; then
        NODE_IPS+=("$node_ip")
        echo -e "${GREEN}  ✓ Added: $node_ip${NC}"
    fi
done

if [ ${#NODE_IPS[@]} -eq 0 ]; then
    echo -e "${RED}Error: No node IPs provided${NC}"
    exit 1
fi

echo ""
echo -e "${GREEN}Node IPs to configure:${NC}"
for ip in "${NODE_IPS[@]}"; do
    echo "  - $ip"
done
echo ""

# Ask for SSH user
read -p "SSH username for nodes (default: ubuntu): " ssh_user
ssh_user=${ssh_user:-ubuntu}

echo ""
echo -e "${YELLOW}Step 2: Configure Registry on Each Node${NC}"
echo ""

# Function to configure a node
configure_node() {
    local node_ip=$1
    local ssh_user=$2

    echo -e "${BLUE}Configuring node: $node_ip${NC}"

    # Check if we can SSH to the node
    if ! ssh -q -o ConnectTimeout=5 -o BatchMode=yes ${ssh_user}@${node_ip} exit; then
        echo -e "${YELLOW}  ⚠ Cannot SSH with key-based auth to ${ssh_user}@${node_ip}${NC}"
        echo -e "${YELLOW}  You may need to enter password...${NC}"
    fi

    # Create the directory
    echo "  Creating /etc/rancher/k3s directory..."
    ssh ${ssh_user}@${node_ip} "sudo mkdir -p /etc/rancher/k3s" || {
        echo -e "${RED}  ✗ Failed to create directory${NC}"
        return 1
    }

    # Copy the registries.yaml file
    echo "  Copying registries.yaml..."
    scp "$SCRIPT_DIR/registries.yaml.example" ${ssh_user}@${node_ip}:/tmp/registries.yaml || {
        echo -e "${RED}  ✗ Failed to copy file${NC}"
        return 1
    }

    # Move to correct location with sudo
    ssh ${ssh_user}@${node_ip} "sudo mv /tmp/registries.yaml /etc/rancher/k3s/registries.yaml" || {
        echo -e "${RED}  ✗ Failed to move file${NC}"
        return 1
    }

    # Set correct permissions
    ssh ${ssh_user}@${node_ip} "sudo chmod 644 /etc/rancher/k3s/registries.yaml" || {
        echo -e "${YELLOW}  ⚠ Warning: Could not set permissions${NC}"
    }

    # Verify file exists
    echo "  Verifying configuration..."
    if ssh ${ssh_user}@${node_ip} "sudo test -f /etc/rancher/k3s/registries.yaml"; then
        echo -e "${GREEN}  ✓ Configuration file installed${NC}"
    else
        echo -e "${RED}  ✗ Configuration file not found after installation${NC}"
        return 1
    fi

    # Detect if this is a server or agent node
    echo "  Detecting node type..."
    if ssh ${ssh_user}@${node_ip} "sudo systemctl list-units --full --all | grep -q k3s.service"; then
        NODE_TYPE="server"
        SERVICE_NAME="k3s"
    elif ssh ${ssh_user}@${node_ip} "sudo systemctl list-units --full --all | grep -q k3s-agent.service"; then
        NODE_TYPE="agent"
        SERVICE_NAME="k3s-agent"
    else
        echo -e "${YELLOW}  ⚠ Could not detect node type, assuming agent${NC}"
        NODE_TYPE="agent"
        SERVICE_NAME="k3s-agent"
    fi

    echo -e "  Node type: ${BLUE}${NODE_TYPE}${NC}"

    # Restart the service
    echo "  Restarting ${SERVICE_NAME} service..."
    if ssh ${ssh_user}@${node_ip} "sudo systemctl restart ${SERVICE_NAME}"; then
        echo -e "${GREEN}  ✓ Service restarted successfully${NC}"
    else
        echo -e "${RED}  ✗ Failed to restart service${NC}"
        echo -e "${YELLOW}  You may need to restart manually:${NC}"
        echo -e "${YELLOW}    ssh ${ssh_user}@${node_ip} 'sudo systemctl restart ${SERVICE_NAME}'${NC}"
        return 1
    fi

    # Wait a moment for service to stabilize
    sleep 2

    # Check service status
    echo "  Checking service status..."
    if ssh ${ssh_user}@${node_ip} "sudo systemctl is-active --quiet ${SERVICE_NAME}"; then
        echo -e "${GREEN}  ✓ Service is running${NC}"
    else
        echo -e "${RED}  ✗ Service is not running!${NC}"
        echo -e "${YELLOW}  Check logs with: ssh ${ssh_user}@${node_ip} 'sudo journalctl -u ${SERVICE_NAME} -n 50'${NC}"
        return 1
    fi

    # Test registry access (with patience for large image)
    echo "  Testing registry access..."
    echo -e "  ${BLUE}Note: Image is ~1-2GB, this may take 1-3 minutes on first pull${NC}"
    if ssh ${ssh_user}@${node_ip} "timeout 300 sudo k3s crictl pull 192.168.1.208:3002/jason/socktop-webterm:0.2.0 2>&1" | grep -q "Image is up to date\|Successfully pulled"; then
        echo -e "${GREEN}  ✓ Successfully pulled image from registry!${NC}"
    else
        echo -e "${YELLOW}  ⚠ Could not confirm image pull (may already be cached or need credentials)${NC}"
        echo -e "${YELLOW}  You can verify manually: ssh ${ssh_user}@${node_ip} 'sudo k3s crictl images | grep socktop'${NC}"
    fi

    echo -e "${GREEN}✓ Node $node_ip configured successfully!${NC}"
    echo ""
    return 0
}

# Configure each node
FAILED_NODES=()
for node_ip in "${NODE_IPS[@]}"; do
    if ! configure_node "$node_ip" "$ssh_user"; then
        FAILED_NODES+=("$node_ip")
    fi
done

echo ""
echo -e "${GREEN}=== Configuration Summary ===${NC}"
echo ""

if [ ${#FAILED_NODES[@]} -eq 0 ]; then
    echo -e "${GREEN}✓ All nodes configured successfully!${NC}"
    echo ""
    echo "Your k3s cluster is now configured to pull images from:"
    echo -e "  ${BLUE}192.168.1.208:3002${NC}"
    echo ""
    echo "You can now deploy Socktop WebTerm with:"
    echo -e "  ${YELLOW}cd kubernetes${NC}"
    echo -e "  ${YELLOW}./deploy.sh${NC}"
else
    echo -e "${RED}✗ Some nodes failed to configure:${NC}"
    for node in "${FAILED_NODES[@]}"; do
        echo -e "  ${RED}- $node${NC}"
    done
    echo ""
    echo "Please configure these nodes manually:"
    echo ""
    echo "1. SSH to the node:"
    echo -e "   ${YELLOW}ssh ${ssh_user}@<node-ip>${NC}"
    echo ""
    echo "2. Create the directory:"
    echo -e "   ${YELLOW}sudo mkdir -p /etc/rancher/k3s${NC}"
    echo ""
    echo "3. Copy the registries.yaml file:"
    echo -e "   ${YELLOW}scp registries.yaml.example ${ssh_user}@<node-ip>:/tmp/registries.yaml${NC}"
    echo -e "   ${YELLOW}ssh ${ssh_user}@<node-ip> 'sudo mv /tmp/registries.yaml /etc/rancher/k3s/registries.yaml'${NC}"
    echo ""
    echo "4. Restart k3s:"
    echo -e "   ${YELLOW}sudo systemctl restart k3s${NC}        # on server nodes"
    echo -e "   ${YELLOW}sudo systemctl restart k3s-agent${NC}  # on agent nodes"
fi

echo ""
echo -e "${GREEN}Done!${NC}"