#!/bin/bash
# Mullvad Browser Container Manager
# Usage: mullvad-browser.sh [create|start|stop|launch|status|destroy|shell]

set -e

CONTAINER="mullvad-browser"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SPEC_FILE="$SCRIPT_DIR/../specs/mullvad-browser.yaml"
DOWNLOADS_DIR="$HOME/Downloads/mullvad-browser"
WG_CONFIG="$HOME/containers/mullvad-configs/dk-cph-wg-001.conf"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

log() { echo -e "${GREEN}[mullvad-browser]${NC} $1"; }
warn() { echo -e "${YELLOW}[mullvad-browser]${NC} $1"; }
error() { echo -e "${RED}[mullvad-browser]${NC} $1"; }

container_exists() {
    sudo machinectl list-images 2>/dev/null | grep -q "^$CONTAINER "
}

container_running() {
    sudo machinectl list 2>/dev/null | grep -q "^$CONTAINER "
}

wait_for_boot() {
    log "Waiting for container to boot..."
    for i in {1..30}; do
        if sudo machinectl shell "$CONTAINER" /bin/true 2>/dev/null; then
            return 0
        fi
        sleep 1
    done
    error "Container failed to boot"
    return 1
}

setup_network() {
    log "Configuring network..."
    sudo machinectl shell "$CONTAINER" /bin/bash -c '
        # Create network config if not exists
        if [ ! -f /etc/systemd/network/80-container-host0.network ]; then
            cat > /etc/systemd/network/80-container-host0.network << EOF
[Match]
Name=host0

[Network]
DHCP=yes
EOF
        fi
        systemctl enable --now systemd-networkd 2>/dev/null || true
    ' 2>/dev/null

    # Wait for IP
    for i in {1..30}; do
        IP=$(sudo machinectl show "$CONTAINER" -p Addresses --value 2>/dev/null | grep -oE '192\.[0-9]+\.[0-9]+\.[0-9]+' | head -1)
        if [ -n "$IP" ]; then
            log "Container IP: $IP"
            return 0
        fi
        sleep 1
    done
    warn "Could not get container IP"
}

setup_wireguard() {
    log "Setting up WireGuard with Mullvad..."
    
    # Copy WireGuard config to container
    sudo machinectl copy-to "$CONTAINER" "$WG_CONFIG" /etc/wireguard/mullvad.conf
    
    sudo machinectl shell "$CONTAINER" /bin/bash -c '
        # Install WireGuard if not present
        if ! command -v wg &>/dev/null; then
            apt-get update
            apt-get install -y wireguard-tools iptables curl iproute2
        fi

        # Extract VPN server endpoint for kill switch
        VPN_ENDPOINT=$(grep "^Endpoint" /etc/wireguard/mullvad.conf | cut -d= -f2 | tr -d " " | cut -d: -f1)
        VPN_PORT=$(grep "^Endpoint" /etc/wireguard/mullvad.conf | cut -d= -f2 | tr -d " " | cut -d: -f2)

        echo "VPN Endpoint: $VPN_ENDPOINT:$VPN_PORT"

        # Bring up WireGuard interface
        wg-quick down mullvad 2>/dev/null || true
        wg-quick up mullvad

        # Configure kill switch iptables rules
        iptables -F
        iptables -X LEAK_CHECK 2>/dev/null || true
        iptables -N LEAK_CHECK

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

        # Allow traffic to WireGuard interface
        iptables -A LEAK_CHECK -o mullvad -j ACCEPT

        # Allow traffic to VPN server (for handshake)
        iptables -A LEAK_CHECK -d $VPN_ENDPOINT -p udp --dport $VPN_PORT -j ACCEPT

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

        # Allow local network (for container networking)
        iptables -A LEAK_CHECK -d 10.0.0.0/8 -j ACCEPT
        iptables -A LEAK_CHECK -d 172.16.0.0/12 -j ACCEPT
        iptables -A LEAK_CHECK -d 192.168.0.0/16 -j ACCEPT

        # Log and drop everything else
        iptables -A LEAK_CHECK -j LOG --log-prefix "MULLVAD-LEAK: " --log-level 4
        iptables -A LEAK_CHECK -j DROP

        # Apply chain to OUTPUT
        iptables -C OUTPUT -j LEAK_CHECK 2>/dev/null || iptables -I OUTPUT 1 -j LEAK_CHECK

        # Block IPv6 to prevent leaks
        ip6tables -P OUTPUT DROP 2>/dev/null || true
        ip6tables -A OUTPUT -o lo -j ACCEPT 2>/dev/null || true

        # Save iptables rules
        mkdir -p /etc/iptables
        iptables-save > /etc/iptables/rules.v4
        ip6tables-save > /etc/iptables/rules.v6 2>/dev/null || true

        # Create restore service
        cat > /etc/systemd/system/iptables-restore.service << EOF
[Unit]
Description=Restore iptables rules
Before=network-pre.target
Wants=network-pre.target

[Service]
Type=oneshot
ExecStart=/sbin/iptables-restore /etc/iptables/rules.v4
ExecStart=/sbin/ip6tables-restore /etc/iptables/rules.v6
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
EOF

        # Create WireGuard auto-start service
        cat > /etc/systemd/system/wg-mullvad.service << EOF
[Unit]
Description=WireGuard Mullvad VPN
After=network-online.target
Wants=network-online.target

[Service]
Type=oneshot
RemainAfterExit=yes
ExecStart=/usr/bin/wg-quick up mullvad
ExecStop=/usr/bin/wg-quick down mullvad

[Install]
WantedBy=multi-user.target
EOF

        systemctl daemon-reload
        systemctl enable iptables-restore
        systemctl enable wg-mullvad
    ' 2>/dev/null

    log "Verifying VPN connection..."
    for i in {1..30}; do
        VPN_IP=$(sudo machinectl shell "$CONTAINER" /bin/bash -c 'curl -s --connect-timeout 5 https://am.i.mullvad.net/ip 2>/dev/null' 2>/dev/null | tr -d '\r')
        if [ -n "$VPN_IP" ] && [ "$VPN_IP" != "$(curl -s ifconfig.me)" ]; then
            log "VPN connected! Mullvad IP: $VPN_IP"
            return 0
        fi
        sleep 2
    done
    warn "VPN connection timeout - may still be connecting"
}

setup_firefox() {
    log "Installing Firefox and graphics libraries..."
    sudo machinectl shell "$CONTAINER" /bin/bash -c '
        if ! command -v firefox-esr &>/dev/null; then
            apt-get update
            apt-get install -y firefox-esr libegl1 libgl1 libgbm1
        fi

        # Create browser user with UID 1000
        if ! id -u browser 2>/dev/null; then
            useradd -u 1000 -m -s /bin/bash browser
        fi

        # Create Firefox profile directory
        PROFILE_DIR="/home/browser/.mozilla/firefox/mullvad.default"
        mkdir -p "$PROFILE_DIR/chrome"

        # Create profiles.ini
        cat > /home/browser/.mozilla/firefox/profiles.ini << PROFILES
[Profile0]
Name=default
IsRelative=1
Path=mullvad.default
Default=1

[General]
StartWithLastProfile=1
PROFILES

        # Create userChrome.css for Mullvad green themed title bar
        cat > "$PROFILE_DIR/chrome/userChrome.css" << CHROME
/* Mullvad Browser - Green Theme Title Bar */
@namespace url("http://www.mozilla.org/keymaster/gatekeeper/there.is.only.xul");

/* Main window title bar - Mullvad green */
#navigator-toolbox {
    background: linear-gradient(to bottom, #44AD4D 0%, #2D8B36 100%) !important;
}

/* Tab bar background */
#titlebar, #TabsToolbar {
    background: #2D8B36 !important;
}

/* Navigation bar */
#nav-bar {
    background: linear-gradient(to bottom, #246B2C 0%, #1A4F20 100%) !important;
    border-bottom: 2px solid #5AC463 !important;
}

/* URL bar - light background with dark text for readability */
#urlbar-background {
    background-color: #d5f0d8 !important;
    border: 2px solid #5AC463 !important;
}

#urlbar {
    background-color: #d5f0d8 !important;
    color: #0a2b0d !important;
}

#urlbar * {
    color: #0a2b0d !important;
}

#urlbar-input {
    color: #0a2b0d !important;
}

#urlbar-input-container {
    color: #0a2b0d !important;
}

#identity-box {
    color: #0a2b0d !important;
}

/* Tab styling - unselected tabs */
.tabbrowser-tab {
    background: #1A4F20 !important;
}

.tabbrowser-tab .tab-background {
    background: #1A4F20 !important;
}

/* Selected tab - lighter green with dark text */
.tabbrowser-tab[selected="true"] {
    background: #5AC463 !important;
}

.tabbrowser-tab[selected="true"] .tab-background {
    background: #5AC463 !important;
}

/* Tab text - light for unselected, dark for selected */
.tabbrowser-tab .tab-label {
    color: #b8e6bb !important;
}

.tabbrowser-tab[selected="true"] .tab-label {
    color: #0a2b0d !important;
    font-weight: bold !important;
}

/* Tab close button */
.tabbrowser-tab .tab-close-button {
    color: #b8e6bb !important;
}

.tabbrowser-tab[selected="true"] .tab-close-button {
    color: #0a2b0d !important;
}

/* Toolbar buttons */
#nav-bar toolbarbutton {
    color: #b8e6bb !important;
}

/* Bookmarks bar */
#PersonalToolbar {
    background: #1A4F20 !important;
}

/* Menu bar */
#toolbar-menubar {
    background: #2D8B36 !important;
    color: white !important;
}
CHROME

        # Enable userChrome.css in Firefox preferences
        cat > "$PROFILE_DIR/user.js" << USERJS
// Enable userChrome.css
user_pref("toolkit.legacyUserProfileCustomizations.stylesheets", true);

// Privacy settings
user_pref("privacy.resistFingerprinting", true);
user_pref("privacy.trackingprotection.enabled", true);
user_pref("media.peerconnection.enabled", false);
user_pref("geo.enabled", false);
user_pref("dom.battery.enabled", false);

// Disable telemetry
user_pref("toolkit.telemetry.enabled", false);
user_pref("datareporting.healthreport.uploadEnabled", false);

// Homepage - Mullvad connection check
user_pref("browser.startup.homepage", "https://mullvad.net/en/check");

// Show title bar (for the colored theme)
user_pref("browser.tabs.inTitlebar", 0);

// Disable DNS over HTTPS to prevent DNS leaks
user_pref("network.trr.mode", 5);
user_pref("network.trr.uri", "");
user_pref("doh-rollout.disable-heuristics", true);
USERJS

        chown -R browser:browser /home/browser/.mozilla
        echo "Firefox profile configured with Mullvad green theme"
    ' 2>/dev/null
}

create_nspawn_config() {
    log "Creating nspawn configuration..."
    sudo tee /etc/systemd/nspawn/$CONTAINER.nspawn > /dev/null << EOF
[Exec]
Boot=yes
PrivateUsers=no
Environment=DISPLAY=:0
Environment=WAYLAND_DISPLAY=wayland-0
Environment=XDG_RUNTIME_DIR=/tmp
Environment=MOZ_ENABLE_WAYLAND=1
Environment=MOZ_WEBRENDER=1
Environment=MOZ_DISABLE_RTC=1

[Network]
VirtualEthernet=yes
Bridge=br0

[Files]
Bind=/tmp/.X11-unix:/tmp/.X11-unix
Bind=/run/user/1000/wayland-0:/tmp/wayland-0
Bind=$DOWNLOADS_DIR:/home/browser/Downloads
EOF
}

cmd_create() {
    if container_exists; then
        error "Container already exists. Use 'destroy' first or 'start' to run it."
        exit 1
    fi

    if [ ! -f "$WG_CONFIG" ]; then
        error "WireGuard config not found: $WG_CONFIG"
        exit 1
    fi

    # Create downloads directory
    mkdir -p "$DOWNLOADS_DIR"

    log "Creating container from spec..."

    # Find gaol binary
    GAOL_BIN=$(which gaol)
    if [ -z "$GAOL_BIN" ]; then
        error "gaol not found in PATH"
        exit 1
    fi

    sudo "$GAOL_BIN" apply "$SPEC_FILE"

    # Update nspawn config with correct settings
    create_nspawn_config

    log "Starting container..."
    sudo machinectl start "$CONTAINER"
    wait_for_boot

    setup_network
    setup_wireguard
    setup_firefox

    log "Container created and configured!"
    log "Run '$0 launch' to start Firefox"
}

cmd_start() {
    if ! container_exists; then
        error "Container does not exist. Run 'create' first."
        exit 1
    fi

    if container_running; then
        warn "Container already running"
        return 0
    fi

    log "Starting container..."
    sudo machinectl start "$CONTAINER"
    wait_for_boot
    setup_network

    # Restore iptables and start WireGuard
    sudo machinectl shell "$CONTAINER" /bin/bash -c '
        /sbin/iptables-restore /etc/iptables/rules.v4 2>/dev/null || true
        /sbin/ip6tables-restore /etc/iptables/rules.v6 2>/dev/null || true
        wg-quick up mullvad 2>/dev/null || true
    ' 2>/dev/null

    log "Waiting for VPN..."
    for i in {1..30}; do
        VPN_IP=$(sudo machinectl shell "$CONTAINER" /bin/bash -c 'curl -s --connect-timeout 5 https://am.i.mullvad.net/ip 2>/dev/null' 2>/dev/null | tr -d '\r')
        if [ -n "$VPN_IP" ]; then
            log "Container started and VPN connected! Mullvad IP: $VPN_IP"
            return 0
        fi
        sleep 2
    done
    warn "VPN may still be connecting"
}

cmd_stop() {
    if ! container_running; then
        warn "Container not running"
        return 0
    fi

    log "Stopping container..."
    sudo machinectl stop "$CONTAINER"
    log "Container stopped"
}

cmd_launch() {
    if ! container_running; then
        log "Container not running, starting..."
        cmd_start
    fi

    log "Launching Firefox with Mullvad theme..."
    sudo machinectl shell "$CONTAINER" /bin/bash -c '
        # Kill existing Firefox if any
        pkill -9 firefox-esr 2>/dev/null || true
        sleep 1

        # Fix dconf permissions
        mkdir -p /tmp/dconf
        chmod 777 /tmp/dconf

        # Launch Firefox as browser user (UID 1000 matches host for Wayland socket access)
        runuser -u browser -- bash -c "
            export WAYLAND_DISPLAY=wayland-0
            export XDG_RUNTIME_DIR=/tmp
            export XDG_CACHE_HOME=/home/browser/.cache
            export XDG_CONFIG_HOME=/home/browser/.config
            export MOZ_ENABLE_WAYLAND=1
            export NO_AT_BRIDGE=1
            export LIBGL_ALWAYS_SOFTWARE=1
            export HOME=/home/browser
            mkdir -p /home/browser/.cache /home/browser/.config
            setsid firefox-esr --profile /home/browser/.mozilla/firefox/mullvad.default https://mullvad.net/en/check </dev/null >/dev/null 2>&1 &
        "
        echo "Firefox launched with Mullvad profile"
    ' 2>/dev/null

    log "Firefox launched with green Mullvad theme - check your desktop"
}

cmd_status() {
    echo "=== Container Status ==="
    if container_exists; then
        echo "Image: EXISTS"
    else
        echo "Image: NOT FOUND"
        return
    fi

    if container_running; then
        echo "State: RUNNING"
        IP=$(sudo machinectl show "$CONTAINER" -p Addresses --value 2>/dev/null | head -1)
        echo "Container IP: $IP"

        # Check WireGuard
        WG_STATUS=$(sudo machinectl shell "$CONTAINER" /bin/bash -c 'wg show mullvad 2>/dev/null | head -1' 2>/dev/null)
        if [ -n "$WG_STATUS" ]; then
            echo "WireGuard: CONNECTED"
            VPN_IP=$(sudo machinectl shell "$CONTAINER" /bin/bash -c 'curl -s --connect-timeout 5 https://am.i.mullvad.net/ip 2>/dev/null' 2>/dev/null | tr -d '\r')
            if [ -n "$VPN_IP" ]; then
                echo "Mullvad IP: $VPN_IP"
            fi
            
            # Check Mullvad connection
            MULLVAD_CHECK=$(sudo machinectl shell "$CONTAINER" /bin/bash -c 'curl -s https://am.i.mullvad.net/connected 2>/dev/null' 2>/dev/null | tr -d '\r')
            echo "Mullvad Status: $MULLVAD_CHECK"
        else
            echo "WireGuard: DISCONNECTED"
        fi

        # Check Firefox
        if sudo machinectl shell "$CONTAINER" /bin/bash -c 'pgrep firefox-esr' 2>/dev/null | grep -q .; then
            echo "Firefox: RUNNING"
        else
            echo "Firefox: NOT RUNNING"
        fi
    else
        echo "State: STOPPED"
    fi
}

cmd_destroy() {
    if container_running; then
        log "Stopping container..."
        sudo machinectl stop "$CONTAINER" 2>/dev/null || true
        sleep 2
    fi

    if container_exists; then
        log "Removing container..."
        sudo machinectl remove "$CONTAINER" 2>/dev/null || true
        sudo rm -f /etc/systemd/nspawn/$CONTAINER.nspawn
        log "Container destroyed"
    else
        warn "Container does not exist"
    fi
}

cmd_shell() {
    if ! container_running; then
        error "Container not running. Run 'start' first."
        exit 1
    fi

    sudo machinectl shell "$CONTAINER"
}

cmd_help() {
    cat << EOF
Mullvad Browser Container Manager

Usage: $0 <command>

Commands:
  create   Create and configure the Mullvad browser container
  start    Start the container (if stopped)
  stop     Stop the container
  launch   Launch Firefox in the container (starts if needed)
  status   Show container and VPN status
  destroy  Remove the container completely
  shell    Open a shell in the container
  help     Show this help message

The container routes ALL traffic through Mullvad VPN via WireGuard.
Server: Denmark (Copenhagen) - dk-cph-wg-001

Downloads are saved to: $DOWNLOADS_DIR
EOF
}

# Main
case "${1:-help}" in
    create)  cmd_create ;;
    start)   cmd_start ;;
    stop)    cmd_stop ;;
    launch)  cmd_launch ;;
    status)  cmd_status ;;
    destroy) cmd_destroy ;;
    shell)   cmd_shell ;;
    help|--help|-h) cmd_help ;;
    *)
        error "Unknown command: $1"
        cmd_help
        exit 1
        ;;
esac
