#!/usr/bin/env bash
set -euo pipefail

# Resolve package directory (follows symlinks from global bin).
# On macOS 12.3+ readlink supports -f; fallback to python3 for older systems / safety.
if readlink -f "$0" >/dev/null 2>&1; then
    SELF="$(readlink -f "$0")"
else
    SELF="$(python3 -c 'import os,sys; print(os.path.realpath(sys.argv[1]))' "$0")"
fi
BIN_DIR="$(dirname "$SELF")"
PACKAGE_DIR="$(dirname "$BIN_DIR")"
LIB_DIR="$PACKAGE_DIR/lib"

# shellcheck source=../lib/bash-helpers.sh
source "$LIB_DIR/bash-helpers.sh"

WAYBAR_CONFIG="$HOME/.config/waybar/config.jsonc"
WAYBAR_STYLE="$HOME/.config/waybar/style.css"

usage() {
    echo "ai-gauge — AI coding assistant usage monitor"
    echo ""
    echo "Usage: ai-gauge <command>"
    echo ""
    echo "Commands:"
    echo "  setup       Install systemd service, patch waybar config + CSS"
    echo "  uninstall   Remove service, waybar config, CSS, runtime state"
    echo "  status      Show service status"
    echo "  version     Show version"
    echo ""
    echo "Configuration:"
    echo "  ai-gauge-config get              Show current config"
    echo "  ai-gauge-config set <key> <val>  Set config value"
}

cmd_version() {
    local pkg="$PACKAGE_DIR/package.json"
    if [[ -f "$pkg" ]] && command -v jq &>/dev/null; then
        jq -r '.version' "$pkg"
    else
        echo "unknown"
    fi
}

cmd_status() {
    if is_macos; then
        local UID_VAL
        UID_VAL=$(id -u)
        
        echo "=== ai-gauge status (macOS) ==="
        echo ""
        echo "Server agent (com.ai-gauge.server):"
        launchctl print "gui/$UID_VAL/com.ai-gauge.server" 2>/dev/null | grep -E "state|pid|last exit" || echo "  (not loaded)"
        echo ""
        echo "Menubar agent (com.ai-gauge.menubar):"
        launchctl print "gui/$UID_VAL/com.ai-gauge.menubar" 2>/dev/null | grep -E "state|pid|last exit" || echo "  (not loaded)"
        echo ""
        echo "State file:"
        jq . "${TMPDIR:-/tmp}/ai-gauge/usage.json" 2>/dev/null || echo "  (no state)"
        echo ""
        echo "Recent server logs:"
        tail -n 5 "$HOME/Library/Logs/ai-gauge/server.log" 2>/dev/null || echo "  (no logs)"
        return 0
    fi
    
    systemctl --user status ai-gauge-server --no-pager 2>&1 || true
}

# Test harness: skip side-effecting steps in dry-run mode
setup_skip_or_run() {
    local desc="$1"; shift
    if [[ "${AIGAUGE_SETUP_DRY_RUN:-0}" == "1" ]]; then
        printf '[setup-dry-run] skip: %s\n' "$desc" >&2
        return 0
    fi
    "$@"
}

cmd_setup_macos() {
    local BUN_PATH
    BUN_PATH=$(command -v bun 2>/dev/null || true)
    if [[ -z "$BUN_PATH" ]]; then
        echo "Error: bun not found in PATH. Install bun first: https://bun.sh"
        exit 1
    fi
    BUN_PATH=$(resolve_path "$BUN_PATH")
    echo "  Found bun at $BUN_PATH"

    local SERVER_PATH
    SERVER_PATH=$(command -v ai-gauge-server 2>/dev/null || true)
    if [[ -z "$SERVER_PATH" ]]; then
        SERVER_PATH="$BIN_DIR/ai-gauge-server"
    fi
    if [[ ! -f "$SERVER_PATH" ]]; then
        echo "Error: ai-gauge-server not found (looked in PATH and $BIN_DIR)"
        exit 1
    fi
    SERVER_PATH=$(resolve_path "$SERVER_PATH")
    echo "  Found ai-gauge-server at $SERVER_PATH"

    local APP_BUNDLE="$BIN_DIR/AIGauge.app"
    local APP_EXEC="$APP_BUNDLE/Contents/MacOS/AIGauge"
    local LOG_DIR="$HOME/Library/Logs/ai-gauge"
    local LAUNCH_AGENTS="$HOME/Library/LaunchAgents"
    local UID_VAL
    UID_VAL=$(id -u)

    if [[ ! -x "$APP_EXEC" ]]; then
        echo "  Menubar app bundle missing; building..."
        if [[ -f "$PACKAGE_DIR/scripts/build-macos-binary.sh" ]]; then
            bash "$PACKAGE_DIR/scripts/build-macos-binary.sh"
        else
            echo "Error: build script not found at $PACKAGE_DIR/scripts/build-macos-binary.sh"
            exit 1
        fi
    fi

    # Create default config if missing (shared with Linux path)
    local CONFIG_DIR="$HOME/.config/ai-gauge"
    local CONFIG_FILE="$CONFIG_DIR/config.json"
    mkdir -p "$CONFIG_DIR"
    ensure_default_config "$CONFIG_FILE" "menubar 'Change token source'"

    _clear_quarantine() {
        xattr -dr com.apple.quarantine "$APP_BUNDLE" 2>/dev/null || true
        echo "  Cleared quarantine attribute on $APP_BUNDLE"
    }

    # Ensure required directories
    _ensure_macos_dirs() {
        mkdir -p "$LOG_DIR"
        mkdir -p "$LAUNCH_AGENTS"
    }
    setup_skip_or_run "macOS log + LaunchAgents directories" _ensure_macos_dirs

    # Generate server plist from template
    local SERVER_PLIST="$LAUNCH_AGENTS/com.ai-gauge.server.plist"
    _write_server_plist() {
        sed -e "s|__BUN_PATH__|$BUN_PATH|g" \
            -e "s|__SERVER_PATH__|$SERVER_PATH|g" \
            -e "s|__LOG_DIR__|$LOG_DIR|g" \
            -e "s|__HOME__|$HOME|g" \
            "$LIB_DIR/ai-gauge-server.plist.template" > "$SERVER_PLIST"
        echo "  Installed $SERVER_PLIST"
    }

    # Generate menubar plist from template
    local MENUBAR_PLIST="$LAUNCH_AGENTS/com.ai-gauge.menubar.plist"
    _write_menubar_plist() {
        sed -e "s|__MENUBAR_APP_EXEC__|$APP_EXEC|g" \
            -e "s|__LOG_DIR__|$LOG_DIR|g" \
            -e "s|__HOME__|$HOME|g" \
            "$LIB_DIR/ai-gauge-menubar.plist.template" > "$MENUBAR_PLIST"
        echo "  Installed $MENUBAR_PLIST"
    }

    _bootstrap_server_agent() {
        launchctl bootout "gui/${UID_VAL}/com.ai-gauge.server" 2>/dev/null || true
        launchctl bootstrap "gui/${UID_VAL}" "$SERVER_PLIST"
        echo "  Bootstrapped com.ai-gauge.server"
    }

    _bootstrap_menubar_agent() {
        launchctl bootout "gui/${UID_VAL}/com.ai-gauge.menubar" 2>/dev/null || true
        launchctl bootstrap "gui/${UID_VAL}" "$MENUBAR_PLIST"
        echo "  Bootstrapped com.ai-gauge.menubar"
    }

    setup_skip_or_run "xattr quarantine removal" _clear_quarantine
    setup_skip_or_run "install server plist" _write_server_plist
    setup_skip_or_run "install menubar plist" _write_menubar_plist

    # Validate plists post-substitution
    setup_skip_or_run "plutil -lint server plist" plutil -lint "$SERVER_PLIST"
    setup_skip_or_run "plutil -lint menubar plist" plutil -lint "$MENUBAR_PLIST"

    # Bootstrap LaunchAgents (idempotent: bootout first, then bootstrap).
    # bootout may return non-zero if agent isn't already loaded — ignore.
    setup_skip_or_run "launchctl bootstrap server" _bootstrap_server_agent
    setup_skip_or_run "launchctl bootstrap menubar" _bootstrap_menubar_agent

    echo "ai-gauge setup complete (macOS)"
    echo "Both LaunchAgents registered. Menubar icon should appear shortly."
}

cmd_setup_linux() {
    local BUN_PATH
    BUN_PATH=$(command -v bun 2>/dev/null || true)
    if [[ -z "$BUN_PATH" ]]; then
        echo "Error: bun not found in PATH. Install bun first: https://bun.sh"
        exit 1
    fi
    BUN_PATH=$(resolve_path "$BUN_PATH")
    echo "  Found bun at $BUN_PATH"

    local SERVER_PATH
    SERVER_PATH=$(command -v ai-gauge-server 2>/dev/null || true)
    if [[ -z "$SERVER_PATH" ]]; then
        SERVER_PATH="$BIN_DIR/ai-gauge-server"
    fi
    if [[ ! -f "$SERVER_PATH" ]]; then
        echo "Error: ai-gauge-server not found (looked in PATH and $BIN_DIR)"
        exit 1
    fi
    SERVER_PATH=$(resolve_path "$SERVER_PATH")
    echo "  Found ai-gauge-server at $SERVER_PATH"

    # Resolve absolute paths for waybar exec/on-click. waybar freezes PATH at
    # start time, so bare command names break silently on bun/pnpm/mise
    # upgrades that move binary symlinks.
    local WAYBAR_BIN_PATH
    WAYBAR_BIN_PATH=$(command -v ai-gauge-waybar 2>/dev/null || true)
    [[ -z "$WAYBAR_BIN_PATH" ]] && WAYBAR_BIN_PATH="$BIN_DIR/ai-gauge-waybar"
    WAYBAR_BIN_PATH=$(resolve_path "$WAYBAR_BIN_PATH")

    local MENU_BIN_PATH
    MENU_BIN_PATH=$(command -v ai-gauge-menu 2>/dev/null || true)
    [[ -z "$MENU_BIN_PATH" ]] && MENU_BIN_PATH="$BIN_DIR/ai-gauge-menu"
    MENU_BIN_PATH=$(resolve_path "$MENU_BIN_PATH")

    # Install systemd service
    local SYSTEMD_DIR="$HOME/.config/systemd/user"
    local SERVICE_SRC="$LIB_DIR/ai-gauge-server.service"
    local SERVICE_DST="$SYSTEMD_DIR/ai-gauge-server.service"

    _ensure_systemd_dir() {
        mkdir -p "$SYSTEMD_DIR"
    }

    _install_systemd_service() {
        sed -e "s|__BUN_PATH__|$BUN_PATH|g" \
            -e "s|__SERVER_PATH__|$SERVER_PATH|g" \
            "$SERVICE_SRC" > "$SERVICE_DST"
        echo "  Installed ai-gauge-server.service → $SERVICE_DST"
    }

    setup_skip_or_run "systemd user directory" _ensure_systemd_dir
    setup_skip_or_run "systemd service install" _install_systemd_service

    # Create default config if missing (write before daemon starts so it reads correct tokenSource)
    local CONFIG_DIR="$HOME/.config/ai-gauge"
    local CONFIG_FILE="$CONFIG_DIR/config.json"
    mkdir -p "$CONFIG_DIR"
    ensure_default_config "$CONFIG_FILE" "right-click menu '🔑 Token source'"

    _enable_systemd_service() {
        # Always enable (idempotent — refreshes symlinks if unit content changed).
        systemctl --user enable ai-gauge-server 2>/dev/null || true
        # Force restart so the daemon picks up any changes to ExecStart in the unit file.
        # Using `restart` (not `enable --now`) because `--now` only starts if stopped;
        # a running daemon would keep using the old ExecStart from before daemon-reload.
        # `restart` also handles the cold-start case correctly (starts if not running).
        systemctl --user restart ai-gauge-server 2>/dev/null || true
        echo "  Enabled and (re)started ai-gauge-server systemd service"
    }

    _verify_systemd_service() {
        sleep 2
        if systemctl --user is-active --quiet ai-gauge-server; then
            echo "  ai-gauge-server is active"
        else
            echo "  Warning: ai-gauge-server not yet active (may need a moment to start)"
        fi
    }

    _install_streamdock() {
        local STREAMDOCK_PLUGINS="$HOME/.wine/drive_c/users/$USER/AppData/Roaming/HotSpot/StreamDock/plugins"
        local PLUGIN_SRC="$LIB_DIR/streamdock-plugin"
        if [[ -d "$STREAMDOCK_PLUGINS" ]] && [[ -d "$PLUGIN_SRC" ]]; then
            local PLUGIN_DST="$STREAMDOCK_PLUGINS/com.ai-gauge.streamdock.sdPlugin"
            mkdir -p "$PLUGIN_DST"
            cp -r "$PLUGIN_SRC/"* "$PLUGIN_DST/"
            echo "  Installed StreamDock plugin → $PLUGIN_DST"
        else
            echo "  StreamDock not found (skipped plugin install)"
        fi
    }

    _patch_waybar_config() {
        if [[ ! -f "$WAYBAR_CONFIG" ]]; then
            return
        fi
        local action
        action=$(
            AI_GAUGE_WAYBAR_BIN="$WAYBAR_BIN_PATH" \
            AI_GAUGE_MENU_BIN="$MENU_BIN_PATH" \
            python3 "$LIB_DIR/patch-waybar-config.py" "$WAYBAR_CONFIG"
        )
        case "$action" in
            inserted) echo "  Patched waybar config.jsonc (added module, exec=$WAYBAR_BIN_PATH)" ;;
            migrated) echo "  Patched waybar config.jsonc (migrated paths to $WAYBAR_BIN_PATH)" ;;
            noop)     echo "  Waybar config.jsonc already up-to-date (exec=$WAYBAR_BIN_PATH)" ;;
            *)        echo "  Waybar config.jsonc patched (status: ${action:-unknown})" ;;
        esac
    }

    _patch_waybar_css() {
        if [[ -f "$WAYBAR_STYLE" ]] && ! grep -q '#custom-ai-gauge' "$WAYBAR_STYLE"; then
            cat >> "$WAYBAR_STYLE" << 'CSS'

/* ai-gauge-start */
#custom-ai-gauge {
  min-width: 12px;
  margin-left: 5px;
  margin-right: 0;
  font-size: 11px;
}

#custom-ai-gauge.normal {
  opacity: 0.7;
}

#custom-ai-gauge.warning {
  color: #c5a555;
}

#custom-ai-gauge.update-available {
  color: #f9a825;
}

#custom-ai-gauge.updating {
  color: #29b6f6;
}

#custom-ai-gauge.update-failed {
  color: #ef5350;
}

#custom-ai-gauge.critical {
  color: #a55555;
}

#custom-ai-gauge.waiting {
  opacity: 0.4;
}
/* ai-gauge-end */
CSS
            echo "  Patched waybar style.css"
        else
            echo "  CSS already patched (skipped)"
        fi
    }

    _restart_waybar() {
        if command -v omarchy-restart-waybar &>/dev/null; then
            omarchy-restart-waybar
        else
            # Hard restart (not SIGUSR2): SIGUSR2 only reloads waybar's config,
            # but the process PATH is frozen at start time. If PATH or installed
            # binary paths changed since waybar was started (e.g. bun install
            # path migration after an ai-gauge upgrade), SIGUSR2 won't pick that
            # up — waybar silently fails to find ai-gauge-waybar and the module
            # disappears. A fresh process inherits the current PATH.
            pkill -x waybar 2>/dev/null
            sleep 0.3
            setsid waybar >/dev/null 2>&1 &
            disown 2>/dev/null || true
        fi
        echo "  Restarted waybar"
    }

    setup_skip_or_run "systemctl --user daemon-reload" systemctl --user daemon-reload
    setup_skip_or_run "systemctl --user enable --now ai-gauge-server" _enable_systemd_service

    # Post-install verification (non-fatal)
    setup_skip_or_run "systemctl active check" _verify_systemd_service
    setup_skip_or_run "streamdock plugin install" _install_streamdock
    setup_skip_or_run "waybar config patch" _patch_waybar_config
    setup_skip_or_run "waybar css patch" _patch_waybar_css
    setup_skip_or_run "waybar restart" _restart_waybar

    echo "Done!"
}

cmd_setup() {
    local platform="${AIGAUGE_SETUP_PLATFORM:-}"
    if [[ -n "$platform" ]]; then
        case "$platform" in
            darwin|macos) cmd_setup_macos ;;
            linux) cmd_setup_linux ;;
            *) echo "Invalid AIGAUGE_SETUP_PLATFORM: $platform" >&2; exit 1 ;;
        esac
    elif is_macos; then
        cmd_setup_macos
    else
        cmd_setup_linux
    fi
}

cmd_uninstall() {
    echo "Uninstalling ai-gauge..."

    if is_macos; then
        local UID_VAL
        UID_VAL=$(id -u)
        local LAUNCH_AGENTS="$HOME/Library/LaunchAgents"

        # Bootout agents (graceful if not loaded)
        launchctl bootout "gui/$UID_VAL/com.ai-gauge.server" 2>/dev/null || true
        launchctl bootout "gui/$UID_VAL/com.ai-gauge.menubar" 2>/dev/null || true

        # Remove plist files
        rm -f "$LAUNCH_AGENTS/com.ai-gauge.server.plist"
        rm -f "$LAUNCH_AGENTS/com.ai-gauge.menubar.plist"

        # Clean state dir (macOS TMPDIR)
        rm -rf "${TMPDIR:-/tmp}/ai-gauge"

        # Preserve config: ~/.config/ai-gauge/config.json is NOT deleted

        echo "ai-gauge uninstalled (macOS)"
        return 0
    fi

    # Stop and disable systemd service first
    systemctl --user stop ai-gauge-server 2>/dev/null || true
    systemctl --user disable ai-gauge-server 2>/dev/null || true
    rm -f "$HOME/.config/systemd/user/ai-gauge-server.service"
    systemctl --user daemon-reload 2>/dev/null || true
    echo "  Stopped and removed ai-gauge-server service"

    # Clean waybar config
    if [[ -f "$WAYBAR_CONFIG" ]] && grep -q '"custom/ai-gauge"' "$WAYBAR_CONFIG"; then
        python3 -c "
import sys

path = sys.argv[1]
with open(path) as f:
    text = f.read()

text = text.replace(', \"custom/ai-gauge\"', '', 1)
text = text.replace('\"custom/ai-gauge\", ', '', 1)

lines = text.splitlines(True)
result = []
skip = False
depth = 0
for line in lines:
    if not skip and '\"custom/ai-gauge\"' in line and ':' in line:
        skip = True
        depth = line.count('{') - line.count('}')
        for j in range(len(result) - 1, -1, -1):
            stripped = result[j].rstrip()
            if stripped.endswith(','):
                result[j] = stripped[:-1] + '\n'
                break
            elif stripped:
                break
        continue
    if skip:
        depth += line.count('{') - line.count('}')
        if depth <= 0:
            skip = False
        continue
    result.append(line)

with open(path, 'w') as f:
    f.writelines(result)
" "$WAYBAR_CONFIG"
        echo "  Cleaned waybar config.jsonc"
    else
        echo "  Config not patched (skipped)"
    fi

    # Clean waybar CSS
    if [[ -f "$WAYBAR_STYLE" ]] && grep -q 'ai-gauge-start' "$WAYBAR_STYLE"; then
        sed -i '/\/\* ai-gauge-start \*\//,/\/\* ai-gauge-end \*\//d' "$WAYBAR_STYLE"
        echo "  Cleaned waybar style.css"
    else
        echo "  CSS not patched (skipped)"
    fi

    # Clean runtime state
    rm -rf "${XDG_RUNTIME_DIR:-/tmp}/ai-gauge"
    echo "  Cleaned runtime state"

    # Clean config
    rm -rf "$HOME/.config/ai-gauge"
    echo "  Removed config"

    # Remove StreamDock plugin
    local STREAMDOCK_PLUGINS="$HOME/.wine/drive_c/users/$USER/AppData/Roaming/HotSpot/StreamDock/plugins"
    local PLUGIN_DST="$STREAMDOCK_PLUGINS/com.ai-gauge.streamdock.sdPlugin"
    if [[ -d "$PLUGIN_DST" ]]; then
        rm -rf "$PLUGIN_DST"
        echo "  Removed StreamDock plugin"
    fi

    # Restart waybar (hard restart for consistency — see _restart_waybar in cmd_setup_linux)
    if command -v omarchy-restart-waybar &>/dev/null; then
        omarchy-restart-waybar
    else
        pkill -x waybar 2>/dev/null
        sleep 0.3
        setsid waybar >/dev/null 2>&1 &
        disown 2>/dev/null || true
    fi
    echo "  Restarted waybar"

    echo "Done! Run 'bun remove -g ai-gauge' to remove the package itself."
}

case "${1:-}" in
    setup)    cmd_setup ;;
    uninstall) cmd_uninstall ;;
    status)   cmd_status ;;
    version)  cmd_version ;;
    -v|--version) cmd_version ;;
    -h|--help|"") usage ;;
    *)
        echo "Unknown command: $1"
        usage
        exit 1
        ;;
esac