#!/bin/bash

if [[ ${DEBUG} -gt 0 ]]; then set -x; fi

# Pod network configuration
POD_NS=$2
SUBNET_CIDR=$3
SUBNET_CIDR_SUFFIX="${SUBNET_CIDR#*/}"

# Virtual Ethernet (veth) devices
POD_VETH="pod-veth-$POD_NS"
POD_IP=$4
POD_CIDR="$POD_IP/$SUBNET_CIDR_SUFFIX"

HOST_VETH="host-veth-$POD_NS"
HOST_IP="${POD_IP%.*}.1"
HOST_CIDR="$HOST_IP/24"

# Container 1 (client)
CONTAINER_NAME="curl-client"
CONTAINER_IMAGE="docker.io/alpine/curl:latest"
CONTAINER_COMMAND="/bin/sh"

# Container 2 (server)
CONTAINER_HTTP_SERVER='http_server.py'
CONTAINER_PORT=$5


delete() {
    echo "Deleting..."

    echo "Removing processes in namespace $POD_NS"
    sudo ip netns pids "$POD_NS" | xargs -r sudo kill || true
    sleep 1

    echo "Removing container"
    sudo ctr --namespace "$POD_NS" task kill -s SIGKILL "$CONTAINER_NAME" || true
    sudo ctr --namespace "$POD_NS" container delete "$CONTAINER_NAME" || true
    sudo ctr --namespace "$POD_NS" images delete $CONTAINER_IMAGE || true
    sleep 1
    sudo ctr namespaces rm "$POD_NS" 

    echo "Removing routing rule"
    sudo iptables -t nat -D POSTROUTING -s "$SUBNET_CIDR" ! -o "$HOST_VETH" -j MASQUERADE || true

    echo "Deleting veth interface $HOST_VETH"    
    sudo ip link delete "$HOST_VETH" || true

    echo "Deleting network namespace: $POD_NS"
    sudo ip netns delete "$POD_NS" || true

    echo "Removing file: ${POD_NS}.env"
    rm -f "${POD_NS}.env"

    echo "Delete complete."
    status
}

trap 'echo "Error occurred at line $LINENO"' ERR

# begin create_namespace
create_namespace() {

    echo "Creating network namespace: $POD_NS"
    sudo ip netns del "$POD_NS" 2>/dev/null || true
    sudo ip netns add "$POD_NS"

    echo "Sleeping to enable the namespace to be created"
    sleep 1

    echo "Creating veth pair: $HOST_VETH <-> $POD_VETH"
    sudo ip link add "$HOST_VETH" type veth peer name "$POD_VETH"

    echo "Moving $POD_VETH to namespace $POD_NS"
    sudo ip link set "$POD_VETH" netns "$POD_NS"

    echo "Assigning IP to $HOST_VETH on host: $HOST_CIDR"
    sudo ip addr add "$HOST_CIDR" dev "$HOST_VETH"
    sudo ip link set "$HOST_VETH" up

    echo "Assigning IP to $POD_VETH in namespace: $POD_CIDR"
    sudo ip netns exec "$POD_NS" ip addr add "$POD_CIDR" dev "$POD_VETH"
    sudo ip netns exec "$POD_NS" ip link set "$POD_VETH" up
    sudo ip netns exec "$POD_NS" ip link set lo up

    echo "Adding default route: $HOST_IP"
    sudo ip netns exec $POD_NS ip route add default via $HOST_IP
}
# end create_namespace

setup_routing() {

    echo "Enabling forwarding"
    sudo sysctl -w net.ipv4.ip_forward=1

    echo "Enabling legacy iptables"
    sudo update-alternatives --set iptables /usr/sbin/iptables-legacy
}

start_container_server() {
    
    echo "Starting container server: ($POD_NS, $POD_IP:$CONTAINER_PORT)"

    sudo ip netns exec "$POD_NS" \
        env POD_IP="$POD_IP" CONTAINER_PORT="$CONTAINER_PORT" \
        nohup python3 $CONTAINER_HTTP_SERVER < /dev/null > $POD_NS.$CONTAINER_HTTP_SERVER.log 2>&1 &

    echo "Container created."
}

start_container_client() {

    echo "Pulling container image $CONTAINER_IMAGE"
    sudo ctr --namespace "$POD_NS" image pull "$CONTAINER_IMAGE"

    echo "Starting container $CONTAINER_IMAGE"
    sudo ctr --namespace $POD_NS run --detach \
        --with-ns="network:/var/run/netns/$POD_NS" \
        "$CONTAINER_IMAGE" "$CONTAINER_NAME" "$CONTAINER_COMMAND"

    echo "Container created."
}

# begin test_connection
test_connection() {
    echo "Making REST call from container $CONTAINER_NAME to container $POD_IP:$CONTAINER_PORT"
    sudo ctr --namespace $POD_NS task exec --exec-id 1 "$CONTAINER_NAME" curl -s "127.0.0.1:$CONTAINER_PORT"
}
# end test_connection

status() {

    echo "Container status:"
    NAMESPACES=$(sudo ctr namespaces ls -q | grep '^pod' || true)
    for NS in $NAMESPACES; do
        CONTAINER=$(sudo ctr --namespace $NS container ls | grep "curl-client")
        printf " %-5s  %-15s  %-15s  %-15s\n" "$NS" "$CONTAINER"
    done
    if [ -z "$NAMESPACES" ]; then
        echo " No ctr namespaces found."
    fi    

    echo "Namespace status:"
    INTERFACES=$(ip link show | grep -E "^[0-9]+: *host-veth-" | awk '{print $2}' | tr -d ':' | cut -d'@' -f1)
    for IFACE in $INTERFACES; do
        HOST_IP=$(ip addr show dev "$IFACE" | grep 'inet ' | awk '{print $2}')
        if [ -z "$HOST_IP" ]; then
            HOST_IP="No IP Assigned"
        fi

        POD_NS=$(echo "$IFACE" | sed 's/host-veth-//')
        NS_IP=$(sudo ip netns exec "$POD_NS" ip addr show | grep 'inet ' | grep -v '127.0.0.1' | awk '{print $2}' | head -n 1)

        printf " %-20s  %-15s <-> %s\n" "$IFACE" "$HOST_IP" "$NS_IP"
    
    done
    if [ -z "$INTERFACES" ]; then
        echo " No interfaces starting with 'host-veth-' found."
    fi
}

generate_env_file() {
    local ENV_FILE="${POD_NS}.env"
    echo "Generating environment file: $ENV_FILE"
    
    cat <<EOF > "$ENV_FILE"
# Generated environment variables for $POD_NS
POD_NS="$POD_NS"
SUBNET_CIDR="$SUBNET_CIDR"
SUBNET_CIDR_SUFFIX="$SUBNET_CIDR_SUFFIX"
POD_VETH="$POD_VETH"
POD_IP="$POD_IP"
POD_CIDR="$POD_CIDR"
HOST_VETH="$HOST_VETH"
HOST_IP="$HOST_IP"
HOST_CIDR="$HOST_CIDR"
CONTAINER_NAME="$CONTAINER_NAME"
CONTAINER_IMAGE="$CONTAINER_IMAGE"
CONTAINER_COMMAND="$CONTAINER_COMMAND"
CONTAINER_HTTP_SERVER="$CONTAINER_HTTP_SERVER"
CONTAINER_PORT="$CONTAINER_PORT"
EOF
    echo "File $ENV_FILE created successfully."
}

usage() {
    echo "Usage: $0 <command> [options]"
    echo "Commands:"
    echo "  create    Create a pod using the <namespace> <subnet> <ip_address> <port>"
    echo "             e.g.," 
    echo "              $0 create pod1 192.168.1.0/24 192.168.1.10 8080"
    echo "              $0 create pod2 192.168.1.0/24 192.168.1.11 8081"
    echo "  test     Tests the connection between the container client and the server using curl"
    echo "             e.g.," 
    echo "              $0 test pod1 192.168.1.0/24 192.168.1.10 8080"
    echo "  delete   Removes containers, routing, and namespaces"
    echo "             e.g.," 
    echo "              $0 delete pod1 192.168.1.0/24" 
    echo "  status   Shows the status of the containers"
}

if [[ $# -eq 0 ]]; then
    usage
fi

case $1 in
    create)
        setup_routing
        create_namespace
        start_container_client
        start_container_server
        generate_env_file
        echo "Pod setup complete. Containers are running in the shared namespace."
        ;;
    test)
        test_connection
        ;;
    delete)
        delete
        ;;
    status)
        status
        ;;
    *)
        echo "Invalid option: $1"
        usage
        exit 1
        ;;
esac
exit 0