#!/bin/bash
# Dev Container Management
# Create and manage development containers from templates with overlays
#
# Usage:
#   dev.sh template                    # Create base template (one-time)
#   dev.sh create <name> [overlays]    # Create dev container
#   dev.sh start <name>                # Start container
#   dev.sh stop <name>                 # Stop container
#   dev.sh shell <name>                # Open shell as dev user
#   dev.sh gui <name>                  # Start GUI (labwc on host)
#   dev.sh list                        # List dev containers
#   dev.sh overlays                    # List available overlays
#   dev.sh destroy <name>              # Remove container
#
# Examples:
#   dev.sh create myproject python node        # Python + Node.js
#   dev.sh create tauri-app rust node tauri    # Full Tauri stack
#   dev.sh create backend python               # Python only
#
# Naming: Containers are prefixed with "dev-" (dev domain convention)
#
# Project Binding:
#   dev.sh bind <name> <host_path>:<container_path>
#   dev.sh unbind <name> <container_path>

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
TEMPLATES_DIR="$SCRIPT_DIR/templates"
OVERLAYS_DIR="$SCRIPT_DIR/overlays"
STATE_DIR="${STATE_DIR:-$SCRIPT_DIR/state}"  # Allow override via environment

TEMPLATE_NAME="dev-template"

# Network whitelist for --restrict-network mode
# Package registries, GitHub, DNS
NETWORK_WHITELIST=(
    "1.1.1.1"           # Cloudflare DNS
    "8.8.8.8"           # Google DNS
    "registry.npmjs.org"
    "registry.yarnpkg.com"
    "crates.io"
    "index.crates.io"
    "static.crates.io"
    "pypi.org"
    "files.pythonhosted.org"
    "proxy.golang.org"
    "sum.golang.org"
    "github.com"
    "api.github.com"
    "raw.githubusercontent.com"
    "codeload.github.com"
    "objects.githubusercontent.com"
    "deb.debian.org"
    "security.debian.org"
    "deb.nodesource.com"
    "astral.sh"         # uv installer
    "cli.github.com"    # gh CLI
)

# Colors
CYAN='\033[0;36m'
GREEN='\033[0;32m'
BLUE='\033[0;34m'
RED='\033[0;31m'
ORANGE='\033[0;33m'
NC='\033[0m'

log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
log_success() { echo -e "${GREEN}[OK]${NC} $1"; }
log_dev() { echo -e "${CYAN}[DEV]${NC} $1"; }
log_warn() { echo -e "${ORANGE}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

# Ensure container name has dev- prefix
normalize_name() {
    local name="$1"
    if [[ "$name" != dev-* ]]; then
        echo "dev-$name"
    else
        echo "$name"
    fi
}

container_exists() {
    machinectl show-image "$1" &>/dev/null
}

is_running() {
    machinectl show "$1" --property=State 2>/dev/null | grep -q "State=running"
}

template_exists() {
    container_exists "$TEMPLATE_NAME"
}

# Create base template
cmd_template() {
    log_dev "Creating dev base template..."
    source "$TEMPLATES_DIR/base.sh"
    create_base_template
}

# List available overlays
cmd_overlays() {
    echo "=== Available Overlays ==="
    echo ""
    for overlay in "$OVERLAYS_DIR"/*.sh; do
        local name=$(basename "$overlay" .sh)
        local desc=""
        case "$name" in
            python) desc="Python 3.12+, uv package manager" ;;
            node)   desc="Node.js LTS, npm, pnpm, yarn" ;;
            rust)   desc="Rustup, cargo, rustfmt, clippy" ;;
            go)     desc="Go toolchain, gopls, delve" ;;
            tauri)  desc="WebKit2GTK, Tauri CLI (needs rust+node)" ;;
            claude) desc="Claude Code CLI, ripgrep, fd-find" ;;
            ai-tools) desc="OpenCode + Crush AI coding assistants" ;;
            shell-tools) desc="bat, eza, fd, ripgrep, fzf, zoxide, starship, delta, gh" ;;
            devcontainer) desc="Devcontainer.json settings (env, lifecycle, features)" ;;
        esac
        printf "  ${CYAN}%-10s${NC} %s\n" "$name" "$desc"
    done
    echo ""
    echo "Usage: dev.sh create <name> <overlay1> [overlay2] ..."
}

# Create dev container
cmd_create() {
    local raw_name="$1"
    shift
    local overlays=("$@")

    if [ -z "$raw_name" ]; then
        log_error "Usage: dev.sh create <name> [overlays...]"
        echo ""
        echo "Examples:"
        echo "  dev.sh create myproject python node"
        echo "  dev.sh create tauri-app rust node tauri"
        exit 1
    fi

    local name=$(normalize_name "$raw_name")

    # Check template exists
    if ! template_exists; then
        log_error "Base template not found. Create it first:"
        echo "  dev.sh template"
        exit 1
    fi

    if container_exists "$name"; then
        log_error "Container $name already exists"
        exit 1
    fi

    log_dev "Creating container: $name"
    if [ ${#overlays[@]} -gt 0 ]; then
        log_info "Overlays: ${overlays[*]}"
    fi

    # Clone from template
    log_info "Cloning from template..."
    sudo machinectl clone "$TEMPLATE_NAME" "$name"

    # Create nspawn config with auto-binds for common dotfiles
    sudo tee /etc/systemd/nspawn/$name.nspawn > /dev/null << EOF
[Exec]
Boot=yes
PrivateUsers=no

[Network]
VirtualEthernet=no

[Files]
# Auto-bound dotfiles (read-only for security)
BindReadOnly=$HOME/.ssh:/home/dev/.ssh
BindReadOnly=$HOME/.gitconfig:/home/dev/.gitconfig
# Claude Code config (read-write)
Bind=$HOME/.claude:/home/dev/.claude
Bind=$HOME/.claude.json:/home/dev/.claude.json
# Cache for faster rebuilds
Bind=$HOME/.cache:/home/dev/.cache
# Add project bindings with: dev.sh bind $name /host/path:/container/path
EOF

    # Create .claude dir and file on host if missing
    mkdir -p "$HOME/.claude"
    touch "$HOME/.claude.json"

    # Start container for overlay application
    log_info "Starting container for configuration..."
    sudo machinectl start "$name"
    sleep 5

    # Apply overlays
    for overlay in "${overlays[@]}"; do
        local overlay_script="$OVERLAYS_DIR/$overlay.sh"
        if [ -f "$overlay_script" ]; then
            source "$overlay_script"
            apply_${overlay}_overlay "$name"
        else
            log_warn "Unknown overlay: $overlay (skipping)"
        fi
    done

    # Update prompt to show container name
    sudo machinectl shell "$name" /bin/bash -c "
sed -i 's/\\[dev\\]/[$name]/' /home/dev/.bashrc
"

    # Stop container
    sudo machinectl stop "$name"

    # Save state (overlays applied)
    mkdir -p "$STATE_DIR"
    echo "${overlays[*]}" > "$STATE_DIR/$name.overlays"

    log_success "Container $name created!"
    echo ""
    echo "Next steps:"
    echo "  dev.sh start $name"
    echo "  dev.sh bind $name ~/Projects/myproject:/home/dev/project"
    echo "  dev.sh shell $name"
}

# Start container
cmd_start() {
    local name=$(normalize_name "$1")
    local restrict_network=false

    # Parse flags
    while [[ "$1" == -* ]]; do
        case "$1" in
            --restrict-network|-r)
                restrict_network=true
                shift
                ;;
            *)
                shift
                ;;
        esac
    done

    if [ -z "$name" ] || [ "$name" = "dev-" ]; then
        log_error "Usage: dev.sh start <name> [--restrict-network]"
        exit 1
    fi

    if ! container_exists "$name"; then
        log_error "Container $name does not exist"
        exit 1
    fi

    if is_running "$name"; then
        log_info "Container $name already running"
        return 0
    fi

    log_dev "Starting $name..."
    sudo machinectl start "$name"

    # Apply network restrictions if requested
    if $restrict_network; then
        apply_network_restrictions "$name"
    fi

    log_success "Container $name started"
    if $restrict_network; then
        log_info "Network restricted to package registries + GitHub only"
    fi
}

# Apply network whitelist restrictions
apply_network_restrictions() {
    local name="$1"

    log_info "Applying network restrictions..."

    # Build iptables rules inside container
    sudo machinectl shell "$name" /bin/bash -c '
# Flush existing rules
iptables -F OUTPUT 2>/dev/null || true

# Allow loopback
iptables -A OUTPUT -o lo -j ACCEPT

# Allow established connections
iptables -A OUTPUT -m state --state ESTABLISHED,RELATED -j ACCEPT

# Allow DNS
iptables -A OUTPUT -p udp --dport 53 -j ACCEPT
iptables -A OUTPUT -p tcp --dport 53 -j ACCEPT
'

    # Resolve and whitelist domains
    for entry in "${NETWORK_WHITELIST[@]}"; do
        if [[ "$entry" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
            # Direct IP
            sudo machinectl shell "$name" /bin/bash -c "iptables -A OUTPUT -d $entry -j ACCEPT"
        else
            # Resolve domain to IPs
            local ips=$(dig +short "$entry" 2>/dev/null | grep -E '^[0-9]+\.' || true)
            for ip in $ips; do
                sudo machinectl shell "$name" /bin/bash -c "iptables -A OUTPUT -d $ip -j ACCEPT"
            done
        fi
    done

    # Log and reject everything else
    sudo machinectl shell "$name" /bin/bash -c '
iptables -A OUTPUT -j LOG --log-prefix "DEV-BLOCKED: " --log-level 4
iptables -A OUTPUT -j REJECT --reject-with icmp-net-unreachable
'
}

# Stop container
cmd_stop() {
    local name=$(normalize_name "$1")

    if [ -z "$1" ]; then
        log_error "Usage: dev.sh stop <name>"
        exit 1
    fi

    if ! is_running "$name"; then
        log_info "Container $name not running"
        return 0
    fi

    log_dev "Stopping $name..."
    sudo machinectl stop "$name"

    log_success "Container $name stopped"
}

# Open shell
cmd_shell() {
    local name=$(normalize_name "$1")

    if [ -z "$1" ]; then
        log_error "Usage: dev.sh shell <name>"
        exit 1
    fi

    if ! is_running "$name"; then
        log_error "Container $name not running. Start with: dev.sh start $name"
        exit 1
    fi

    sudo machinectl shell "$name" /bin/su - dev
}

# Start GUI with host-side labwc (trusted trusted border)
cmd_gui() {
    local name=$(normalize_name "$1")
    local config_dir="$SCRIPT_DIR/labwc-config"
    local socket_dir="/run/user/$(id -u)/labwc-$name"
    local socket_path="$socket_dir/wayland-0"
    local container_socket="/tmp/labwc-wayland-0"

    if [ -z "$1" ]; then
        log_error "Usage: dev.sh gui <name>"
        exit 1
    fi

    if ! is_running "$name"; then
        log_info "Starting container..."
        cmd_start "$1"
        sleep 2
    fi

    # Check for labwc config
    if [ ! -d "$config_dir" ]; then
        log_error "labwc config not found at $config_dir"
        exit 1
    fi

    # Kill any existing labwc for this container
    pkill -f "labwc.*labwc-$name" 2>/dev/null || true
    sleep 1

    # Clean up and create socket directory
    rm -rf "$socket_dir" 2>/dev/null
    mkdir -p "$socket_dir"

    # Start labwc on HOST with dev theme (cyan border)
    log_dev "Starting labwc on host for $name..."
    WAYLAND_DISPLAY="/run/user/$(id -u)/${WAYLAND_DISPLAY:-wayland-0}" \
    XDG_RUNTIME_DIR="$socket_dir" \
    WLR_BACKENDS=wayland \
    labwc -C "$config_dir" 2>/dev/null &
    local labwc_pid=$!

    # Wait for socket
    for i in {1..30}; do
        if [ -S "$socket_path" ]; then
            break
        fi
        sleep 0.2
    done

    if [ ! -S "$socket_path" ]; then
        log_error "labwc failed to create socket"
        kill $labwc_pid 2>/dev/null
        exit 1
    fi
    log_info "labwc running (pid $labwc_pid)"

    # Create mount point in container and bind-mount the socket
    sudo machinectl shell "$name" /bin/bash -c "rm -f $container_socket 2>/dev/null; touch $container_socket" 2>/dev/null
    sudo machinectl bind "$name" "$socket_path" "$container_socket"

    # Start terminal in container connecting to host's labwc
    log_info "Starting terminal..."
    sudo machinectl shell "$name" /bin/bash -c "
        su dev -c '
            export WAYLAND_DISPLAY=$container_socket
            export XDG_RUNTIME_DIR=/tmp
            export HOME=/home/dev
            # Try available terminals in order of preference
            if command -v foot &>/dev/null; then
                foot &
            elif command -v alacritty &>/dev/null; then
                alacritty &
            elif command -v weston-terminal &>/dev/null; then
                weston-terminal &
            else
                # Fallback: try to start xterm via XWayland
                xterm &
            fi
        '
    " 2>/dev/null &

    sleep 2
    log_success "GUI started with host-side labwc decoration (trusted trusted border)"
}

# Bind mount a directory
cmd_bind() {
    local name=$(normalize_name "$1")
    local binding="$2"

    if [ -z "$1" ] || [ -z "$binding" ]; then
        log_error "Usage: dev.sh bind <name> <host_path>:<container_path>"
        echo ""
        echo "Example: dev.sh bind myproject ~/Projects/foo:/home/dev/foo"
        exit 1
    fi

    if ! container_exists "$name"; then
        log_error "Container $name does not exist"
        exit 1
    fi

    # Parse binding
    local host_path="${binding%%:*}"
    local container_path="${binding##*:}"

    # Expand ~ in host path
    host_path="${host_path/#\~/$HOME}"

    if [ ! -d "$host_path" ]; then
        log_error "Host path does not exist: $host_path"
        exit 1
    fi

    # Add to nspawn config
    local config="/etc/systemd/nspawn/$name.nspawn"
    if ! grep -q "Bind=$host_path:$container_path" "$config" 2>/dev/null; then
        sudo bash -c "echo 'Bind=$host_path:$container_path' >> $config"
        log_success "Bound $host_path -> $container_path"
        log_info "Restart container to apply: dev.sh stop $name && dev.sh start $name"
    else
        log_info "Binding already exists"
    fi
}

# Remove a bind mount
cmd_unbind() {
    local name=$(normalize_name "$1")
    local container_path="$2"

    if [ -z "$1" ] || [ -z "$container_path" ]; then
        log_error "Usage: dev.sh unbind <name> <container_path>"
        exit 1
    fi

    local config="/etc/systemd/nspawn/$name.nspawn"
    if [ -f "$config" ]; then
        sudo sed -i "\|:$container_path$|d" "$config"
        log_success "Removed binding for $container_path"
        log_info "Restart container to apply"
    fi
}

# List dev containers
cmd_list() {
    echo "=== Dev Containers ==="
    echo ""

    local found=0
    while read -r image; do
        if [[ "$image" == dev-* ]]; then
            found=1
            echo -n "  $image "

            if is_running "$image"; then
                echo -e "${GREEN}running${NC}"
            else
                echo -e "${CYAN}stopped${NC}"
            fi

            # Show overlays
            if [ -f "$STATE_DIR/$image.overlays" ]; then
                local overlays=$(cat "$STATE_DIR/$image.overlays")
                if [ -n "$overlays" ]; then
                    echo "    overlays: $overlays"
                fi
            fi

            # Show bindings
            local config="/etc/systemd/nspawn/$image.nspawn"
            if [ -f "$config" ]; then
                grep "^Bind=" "$config" 2>/dev/null | while read -r line; do
                    echo "    ${line#Bind=}"
                done
            fi
        fi
    done < <(machinectl list-images --no-legend 2>/dev/null | awk '{print $1}')

    if [ $found -eq 0 ]; then
        echo "  (no dev containers)"
        echo ""
        echo "Create one with: dev.sh create <name> [overlays...]"
    fi
    echo ""
}

# Destroy container
cmd_destroy() {
    local name=$(normalize_name "$1")

    if [ -z "$1" ]; then
        log_error "Usage: dev.sh destroy <name>"
        exit 1
    fi

    if [ "$name" = "$TEMPLATE_NAME" ]; then
        log_error "Cannot destroy the template. All containers depend on it."
        exit 1
    fi

    if ! container_exists "$name"; then
        log_error "Container $name does not exist"
        exit 1
    fi

    log_warn "This will permanently remove container: $name"
    read -p "Are you sure? [y/N] " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        log_info "Cancelled"
        exit 0
    fi

    if is_running "$name"; then
        cmd_stop "${name#dev-}"
    fi

    log_info "Removing container..."
    sudo machinectl remove "$name" 2>/dev/null || true
    sudo rm -f "/etc/systemd/nspawn/$name.nspawn"
    rm -f "$STATE_DIR/$name.overlays"

    log_success "Container $name destroyed"
}

# Show status
cmd_status() {
    local name=$(normalize_name "$1")

    if [ -z "$1" ]; then
        cmd_list
        return
    fi

    if ! container_exists "$name"; then
        log_error "Container $name does not exist"
        exit 1
    fi

    echo "=== $name ==="
    echo ""

    echo -n "Status: "
    if is_running "$name"; then
        echo -e "${GREEN}running${NC}"
    else
        echo -e "${CYAN}stopped${NC}"
    fi

    echo -e "Domain: ${CYAN}dev${NC}"

    if [ -f "$STATE_DIR/$name.overlays" ]; then
        local overlays=$(cat "$STATE_DIR/$name.overlays")
        if [ -n "$overlays" ]; then
            echo "Overlays: $overlays"
        fi
    fi

    echo ""
    echo "Bindings:"
    local config="/etc/systemd/nspawn/$name.nspawn"
    if [ -f "$config" ]; then
        grep "^Bind=" "$config" 2>/dev/null | while read -r line; do
            echo "  ${line#Bind=}"
        done
    fi
}

# Main dispatch
case "${1:-help}" in
    template)
        cmd_template
        ;;
    create)
        shift
        cmd_create "$@"
        ;;
    start)
        cmd_start "$2"
        ;;
    stop)
        cmd_stop "$2"
        ;;
    shell)
        cmd_shell "$2"
        ;;
    gui)
        cmd_gui "$2"
        ;;
    bind)
        cmd_bind "$2" "$3"
        ;;
    unbind)
        cmd_unbind "$2" "$3"
        ;;
    list|ls)
        cmd_list
        ;;
    status)
        cmd_status "$2"
        ;;
    overlays)
        cmd_overlays
        ;;
    destroy|rm)
        cmd_destroy "$2"
        ;;
    help|--help|-h)
        echo "Dev Container Management"
        echo ""
        echo "Usage: dev.sh <command> [options]"
        echo ""
        echo "Setup:"
        echo "  template              Create base template (one-time)"
        echo "  overlays              List available overlays"
        echo ""
        echo "Container Management:"
        echo "  create <name> [ovl]   Create container with overlays"
        echo "  start <name>          Start container"
        echo "  stop <name>           Stop container"
        echo "  shell <name>          Open shell as dev user"
        echo "  gui <name>            Start GUI (labwc, trusted border)"
        echo "  list                  List dev containers"
        echo "  status [name]         Show container status"
        echo "  destroy <name>        Remove container"
        echo ""
        echo "Project Binding:"
        echo "  bind <name> <h>:<c>   Bind host path to container path"
        echo "  unbind <name> <c>     Remove binding"
        echo ""
        echo "Examples:"
        echo "  dev.sh template"
        echo "  dev.sh create myproj python node"
        echo "  dev.sh bind myproj ~/Projects/foo:/home/dev/foo"
        echo "  dev.sh start myproj"
        echo "  dev.sh shell myproj"
        echo ""
        echo "Naming: Containers auto-prefixed with 'dev-'"
        ;;
    *)
        log_error "Unknown command: $1"
        echo "Use 'dev.sh help' for usage"
        exit 1
        ;;
esac
