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

set -e

CONTAINER="work-browser"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SPEC_FILE="$SCRIPT_DIR/../specs/work-browser.yaml"
DATA_DIR="$HOME/containers/work-browser"
DOWNLOADS_DIR="$DATA_DIR/downloads"
FIREFOX_DATA="$DATA_DIR/firefox-data"

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

log() { echo -e "${ORANGE}[work-browser]${NC} $1"; }
warn() { echo -e "${YELLOW}[work-browser]${NC} $1"; }
error() { echo -e "${RED}[work-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_firefox() {
    log "Setting up Firefox profile..."

    # Ensure profile directories exist on host
    mkdir -p "$FIREFOX_DATA/firefox/default/chrome"
    mkdir -p "$DOWNLOADS_DIR"

    # Create profiles.ini if it doesn't exist
    if [ ! -f "$FIREFOX_DATA/firefox/profiles.ini" ]; then
        cat > "$FIREFOX_DATA/firefox/profiles.ini" << 'PROFILES'
[Profile0]
Name=default
IsRelative=1
Path=default
Default=1

[General]
StartWithLastProfile=1
PROFILES
    fi

    # Create userChrome.css with orange theme
    cat > "$FIREFOX_DATA/firefox/default/chrome/userChrome.css" << 'CHROME'
/* Work Browser - Orange Theme Title Bar */
@namespace url("http://www.mozilla.org/keymaster/gatekeeper/there.is.only.xul");

/* Main window title bar - Work orange */
#navigator-toolbox {
    background: linear-gradient(to bottom, #b45a14 0%, #8a4510 100%) !important;
}

/* Tab bar background */
#titlebar, #TabsToolbar {
    background: #8a4510 !important;
}

/* Navigation bar */
#nav-bar {
    background: linear-gradient(to bottom, #7a3d0e 0%, #5e2f0b 100%) !important;
    border-bottom: 2px solid #d06a1a !important;
}

/* URL bar */
#urlbar-background {
    background-color: #4a2508 !important;
    border: 1px solid #d06a1a !important;
}

#urlbar-input {
    color: #ffd9b3 !important;
}

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

.tabbrowser-tab .tab-background {
    background: #5e2f0b !important;
}

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

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

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

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

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

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

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

/* Bookmarks bar */
#PersonalToolbar {
    background: #5e2f0b !important;
}

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

    # Create user.js
    cat > "$FIREFOX_DATA/firefox/default/user.js" << 'USERJS'
// Enable userChrome.css
user_pref("toolkit.legacyUserProfileCustomizations.stylesheets", true);

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

    log "Firefox profile configured with orange theme"
}

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

[Network]
Private=no

[Files]
Bind=/run/user/1000/wayland-0:/tmp/wayland-0
Bind=$FIREFOX_DATA/firefox:/home/browser/.mozilla/firefox
Bind=$DOWNLOADS_DIR:/home/browser/Downloads
Bind=/run/user/1000/pipewire-0:/run/user/1000/pipewire-0
Bind=/run/user/1000/pulse:/run/user/1000/pulse
EOF
}

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

    # Setup directories and Firefox profile
    setup_firefox

    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
    create_nspawn_config

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

    # Install Firefox inside container
    sudo machinectl shell "$CONTAINER" /bin/bash -c '
        apt-get update
        apt-get install -y firefox-esr libegl1 libgl1 libgbm1

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

    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
    log "Container started!"
}

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..."
    sudo machinectl shell "$CONTAINER" /bin/bash -c '
        # Ensure browser user exists with UID 1000
        if ! id -u browser 2>/dev/null; then
            useradd -u 1000 -m -s /bin/bash browser 2>/dev/null || true
        fi

        # 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
        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 -P default </dev/null >/dev/null 2>&1 &
        "
        echo "Firefox launched"
    ' 2>/dev/null

    log "Firefox launched with orange 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"

        # 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"
        log "Note: Firefox data preserved in $DATA_DIR"
    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
Work Browser Container Manager

Usage: $0 <command>

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

Firefox data is stored in: $DATA_DIR
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
