#!/usr/bin/env bash
# ──────────────────────────────────────────────────────────────────────
# jupyter-hub — Cross-project Jupyter instance manager
#
# Discovers and manages ALL JupyterLab instances across your projects.
# Prevents port conflicts, finds orphan processes, provides unified
# status and stop-all functionality.
#
# Part of: https://github.com/saymrwulf/JupyterManager
#
# Usage:
#   jupyter-hub status          Show all running Jupyter instances
#   jupyter-hub stop-all        Stop all running Jupyter instances
#   jupyter-hub ports           Show port allocation across 8888-8899
#   jupyter-hub orphans         Find orphan Jupyter processes
#   jupyter-hub kill-orphans    Kill orphan Jupyter processes
#   jupyter-hub which PORT      Show which project owns a port
#   jupyter-hub config          Show current configuration
# ──────────────────────────────────────────────────────────────────────
set -euo pipefail

JHUB_VERSION="1.0.0"
JHUB_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
JHUB_CONFIG_DIR="${XDG_CONFIG_HOME:-$HOME/.config}/jupyter-hub"
JHUB_CONFIG_FILE="$JHUB_CONFIG_DIR/config.sh"

# ── Colours ───────────────────────────────────────────────────────────
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
DIM='\033[2m'
NC='\033[0m'

# ── Default configuration ────────────────────────────────────────────
# Override in $JHUB_CONFIG_FILE
PORT_RANGE_START=8888
PORT_RANGE_END=8899
SCAN_DIRS=()

# ── Load configuration ───────────────────────────────────────────────

_load_config() {
    # Load project-shipped defaults
    if [[ -f "$JHUB_ROOT/config/defaults.sh" ]]; then
        source "$JHUB_ROOT/config/defaults.sh"
    fi

    # Load user overrides
    if [[ -f "$JHUB_CONFIG_FILE" ]]; then
        source "$JHUB_CONFIG_FILE"
    fi

    # Fallback if no scan dirs configured
    if [[ ${#SCAN_DIRS[@]} -eq 0 ]]; then
        SCAN_DIRS=(
            "$HOME/GitClone/ClaudeCodeProjects"
            "$HOME/GitClone/CodexProjects"
        )
    fi
}

# ── Discovery ─────────────────────────────────────────────────────────

_discover_projects() {
    local projects=()
    for scan_dir in "${SCAN_DIRS[@]}"; do
        if [[ -d "$scan_dir" ]]; then
            while IFS= read -r -d '' app_sh; do
                local project_dir
                project_dir="$(dirname "$(dirname "$app_sh")")"
                # Verify it's a Jupyter-capable project
                if [[ -d "$project_dir/notebooks" ]] || grep -q "jupyter" "$project_dir/pyproject.toml" 2>/dev/null; then
                    projects+=("$project_dir")
                fi
            done < <(find "$scan_dir" -maxdepth 3 -path "*/scripts/app.sh" -print0 2>/dev/null)
        fi
    done
    printf '%s\n' "${projects[@]}" | sort -u
}

_project_name() {
    basename "$1"
}

_pid_alive() {
    kill -0 "$1" 2>/dev/null
}

_find_server_info() {
    local project_dir="$1"

    # Strategy 1: .logs/jupyter.pid
    local pid_file="$project_dir/.logs/jupyter.pid"
    if [[ -f "$pid_file" ]]; then
        local pid
        pid=$(cat "$pid_file")
        if [[ -n "$pid" ]] && _pid_alive "$pid"; then
            echo "pid=$pid source=pid_file"
            return 0
        fi
    fi

    # Strategy 2: .run/jupyter.pid
    pid_file="$project_dir/.run/jupyter.pid"
    if [[ -f "$pid_file" ]]; then
        local pid
        pid=$(cat "$pid_file")
        if [[ -n "$pid" ]] && _pid_alive "$pid"; then
            echo "pid=$pid source=run_pid"
            return 0
        fi
    fi

    # Strategy 3: Jupyter runtime JSON files
    for runtime_dir in "$project_dir/.jupyter_runtime" "$project_dir/.run"; do
        if [[ -d "$runtime_dir" ]]; then
            for json_file in "$runtime_dir"/jpserver-*.json; do
                if [[ -f "$json_file" ]]; then
                    local pid port token
                    pid=$(python3 -c "import json; d=json.load(open('$json_file')); print(d.get('pid',''))" 2>/dev/null || true)
                    port=$(python3 -c "import json; d=json.load(open('$json_file')); print(d.get('port',''))" 2>/dev/null || true)
                    token=$(python3 -c "import json; d=json.load(open('$json_file')); print(d.get('token',''))" 2>/dev/null || true)
                    if [[ -n "$pid" ]] && _pid_alive "$pid"; then
                        echo "pid=$pid port=$port token=$token source=runtime_json"
                        return 0
                    fi
                fi
            done
        fi
    done

    return 1
}

_pid_port() {
    local pid="$1"
    lsof -Pan -p "$pid" -i TCP -sTCP:LISTEN 2>/dev/null | awk 'NR>1 {
        split($9, a, ":"); print a[length(a)]
    }' | head -1
}

# ── Commands ──────────────────────────────────────────────────────────

cmd_status() {
    echo ""
    echo -e "  ${BOLD}Jupyter Hub${NC} v${JHUB_VERSION}"
    echo -e "  ${DIM}$(date '+%Y-%m-%d %H:%M:%S')${NC}"
    echo ""

    local found_any=false

    while IFS= read -r project_dir; do
        [[ -z "$project_dir" ]] && continue
        local name
        name=$(_project_name "$project_dir")

        local info
        if info=$(_find_server_info "$project_dir"); then
            found_any=true
            local pid port
            pid=$(echo "$info" | grep -o 'pid=[0-9]*' | cut -d= -f2)
            port=$(_pid_port "$pid" 2>/dev/null || echo "?")
            echo -e "  ${GREEN}●${NC} ${BOLD}$name${NC}"
            echo -e "    PID: $pid  Port: $port"
            echo -e "    Dir: ${DIM}$project_dir${NC}"
            if [[ "$port" != "?" ]]; then
                echo -e "    URL: http://localhost:$port/lab"
            fi
        else
            echo -e "  ${DIM}○ $name${NC}  ${DIM}(stopped)${NC}"
        fi
    done < <(_discover_projects)

    if ! $found_any; then
        echo -e "  ${DIM}No running Jupyter instances found.${NC}"
    fi

    echo ""
}

cmd_ports() {
    echo ""
    echo -e "  ${BOLD}Port Allocation (${PORT_RANGE_START}-${PORT_RANGE_END})${NC}"
    echo ""

    for port in $(seq "$PORT_RANGE_START" "$PORT_RANGE_END"); do
        local pid
        pid=$(lsof -ti :"$port" 2>/dev/null | head -1 || true)
        if [[ -n "$pid" ]]; then
            local cmd_line project_name="unknown"
            cmd_line=$(ps -p "$pid" -o args= 2>/dev/null || true)
            if [[ "$cmd_line" =~ notebook-dir[=\ ]*([^ ]*) ]]; then
                project_name=$(basename "$(dirname "${BASH_REMATCH[1]}")")
            elif [[ "$cmd_line" =~ root_dir[=\ ]*([^ ]*) ]]; then
                project_name=$(basename "${BASH_REMATCH[1]}")
            fi
            echo -e "  ${RED}$port${NC}  pid:$pid  ${BOLD}$project_name${NC}"
            echo -e "        ${DIM}$(echo "$cmd_line" | head -c 100)${NC}"
        else
            echo -e "  ${GREEN}$port${NC}  ${DIM}free${NC}"
        fi
    done
    echo ""
}

cmd_stop_all() {
    echo ""
    local stopped=0

    while IFS= read -r project_dir; do
        [[ -z "$project_dir" ]] && continue
        local name
        name=$(_project_name "$project_dir")

        local info
        if info=$(_find_server_info "$project_dir"); then
            local pid
            pid=$(echo "$info" | grep -o 'pid=[0-9]*' | cut -d= -f2)

            if [[ -x "$project_dir/scripts/app.sh" ]]; then
                echo -e "  Stopping ${BOLD}$name${NC} via app.sh..."
                (cd "$project_dir" && bash scripts/app.sh stop 2>/dev/null) || true
            else
                echo -e "  Stopping ${BOLD}$name${NC} (pid $pid)..."
                kill -TERM "$pid" 2>/dev/null || true
            fi
            stopped=$((stopped + 1))
        fi
    done < <(_discover_projects)

    # Also kill any orphans not claimed by known projects
    local orphan_count=0
    while IFS= read -r pid; do
        [[ -z "$pid" ]] && continue
        # Skip PIDs we already stopped via app.sh
        local already_stopped=false
        while IFS= read -r project_dir; do
            [[ -z "$project_dir" ]] && continue
            for pf in "$project_dir/.logs/jupyter.pid" "$project_dir/.run/jupyter.pid"; do
                if [[ -f "$pf" ]] && [[ "$(cat "$pf" 2>/dev/null)" == "$pid" ]]; then
                    already_stopped=true
                    break 2
                fi
            done
        done < <(_discover_projects)

        if ! $already_stopped && _pid_alive "$pid"; then
            echo -e "  Stopping orphan pid $pid..."
            kill -TERM "$pid" 2>/dev/null || true
            orphan_count=$((orphan_count + 1))
        fi
    done < <(pgrep -f "jupyter.*lab" 2>/dev/null || true)

    local total=$((stopped + orphan_count))
    if (( total > 0 )); then
        echo ""
        echo -e "  ${GREEN}Stopped $total instance(s).${NC}"
    else
        echo -e "  ${DIM}No running instances to stop.${NC}"
    fi
    echo ""
}

cmd_orphans() {
    echo ""
    echo -e "  ${BOLD}Orphan Jupyter Processes${NC}"
    echo -e "  ${DIM}(Jupyter processes not tracked by any project's PID file)${NC}"
    echo ""

    local found=false
    while IFS= read -r line; do
        [[ -z "$line" ]] && continue
        local pid
        pid=$(echo "$line" | awk '{print $2}')

        local claimed=false
        while IFS= read -r project_dir; do
            [[ -z "$project_dir" ]] && continue
            for pid_file in "$project_dir/.logs/jupyter.pid" "$project_dir/.run/jupyter.pid"; do
                if [[ -f "$pid_file" ]] && [[ "$(cat "$pid_file")" == "$pid" ]]; then
                    claimed=true
                    break 2
                fi
            done
            for json_file in "$project_dir/.jupyter_runtime"/jpserver-*.json; do
                if [[ -f "$json_file" ]]; then
                    local json_pid
                    json_pid=$(python3 -c "import json; print(json.load(open('$json_file')).get('pid',''))" 2>/dev/null || true)
                    if [[ "$json_pid" == "$pid" ]]; then
                        claimed=true
                        break 2
                    fi
                fi
            done
        done < <(_discover_projects)

        if ! $claimed; then
            found=true
            local port cmd_short
            port=$(_pid_port "$pid" 2>/dev/null || echo "?")
            cmd_short=$(ps -p "$pid" -o args= 2>/dev/null | head -c 100 || echo "?")
            echo -e "  ${YELLOW}!${NC} PID $pid  Port: $port"
            echo -e "    ${DIM}$cmd_short${NC}"
        fi
    done < <(ps aux | grep "[j]upyter.*lab" || true)

    if ! $found; then
        echo -e "  ${GREEN}No orphan processes found.${NC}"
    else
        echo ""
        echo "  To kill orphans: jupyter-hub kill-orphans"
    fi
    echo ""
}

cmd_kill_orphans() {
    local orphans=()
    while IFS= read -r line; do
        [[ -z "$line" ]] && continue
        local pid
        pid=$(echo "$line" | awk '{print $2}')

        local claimed=false
        while IFS= read -r project_dir; do
            [[ -z "$project_dir" ]] && continue
            for pid_file in "$project_dir/.logs/jupyter.pid" "$project_dir/.run/jupyter.pid"; do
                if [[ -f "$pid_file" ]] && [[ "$(cat "$pid_file")" == "$pid" ]]; then
                    claimed=true
                    break 2
                fi
            done
        done < <(_discover_projects)

        if ! $claimed; then
            orphans+=("$pid")
        fi
    done < <(ps aux | grep "[j]upyter.*lab" || true)

    if [[ ${#orphans[@]} -eq 0 ]]; then
        echo "No orphan Jupyter processes found."
        return 0
    fi

    for pid in "${orphans[@]}"; do
        echo "Killing orphan pid $pid..."
        kill -TERM "$pid" 2>/dev/null || true
    done
    echo "Killed ${#orphans[@]} orphan(s)."
}

cmd_which() {
    local port="${1:-}"
    if [[ -z "$port" ]]; then
        echo "Usage: jupyter-hub which PORT"
        exit 1
    fi

    local pid
    pid=$(lsof -ti :"$port" 2>/dev/null | head -1 || true)
    if [[ -z "$pid" ]]; then
        echo "Port $port is free."
        return 0
    fi

    echo "Port $port is used by pid $pid"
    local cmd_line
    cmd_line=$(ps -p "$pid" -o args= 2>/dev/null || echo "unknown")
    echo "Command: $cmd_line"

    while IFS= read -r project_dir; do
        [[ -z "$project_dir" ]] && continue
        local info
        if info=$(_find_server_info "$project_dir" 2>/dev/null); then
            local proj_pid
            proj_pid=$(echo "$info" | grep -o 'pid=[0-9]*' | cut -d= -f2)
            if [[ "$proj_pid" == "$pid" ]]; then
                echo "Project: $(_project_name "$project_dir")"
                echo "Path: $project_dir"
                return 0
            fi
        fi
    done < <(_discover_projects)

    echo "Project: (not matched to any known project)"
}

cmd_config() {
    echo ""
    echo -e "  ${BOLD}Jupyter Hub Configuration${NC}"
    echo ""
    echo "  Install root:  $JHUB_ROOT"
    echo "  Config dir:    $JHUB_CONFIG_DIR"
    echo "  Config file:   $JHUB_CONFIG_FILE"
    if [[ -f "$JHUB_CONFIG_FILE" ]]; then
        echo "  Config status: loaded"
    else
        echo "  Config status: using defaults"
    fi
    echo "  Port range:    ${PORT_RANGE_START}-${PORT_RANGE_END}"
    echo ""
    echo "  Scan directories:"
    for d in "${SCAN_DIRS[@]}"; do
        if [[ -d "$d" ]]; then
            echo -e "    ${GREEN}✓${NC} $d"
        else
            echo -e "    ${RED}✗${NC} $d (not found)"
        fi
    done
    echo ""
    echo "  Discovered projects:"
    while IFS= read -r project_dir; do
        [[ -z "$project_dir" ]] && continue
        echo "    $(_project_name "$project_dir")  →  $project_dir"
    done < <(_discover_projects)
    echo ""
}

# ── Main ──────────────────────────────────────────────────────────────

_load_config

case "${1:-help}" in
    status)       cmd_status ;;
    ports)        cmd_ports ;;
    stop-all)     cmd_stop_all ;;
    orphans)      cmd_orphans ;;
    kill-orphans) cmd_kill_orphans ;;
    which)        cmd_which "${2:-}" ;;
    config)       cmd_config ;;
    version|--version|-v)
        echo "jupyter-hub v${JHUB_VERSION}"
        ;;
    help|--help|-h)
        echo ""
        echo -e "  ${BOLD}jupyter-hub${NC} v${JHUB_VERSION} — cross-project Jupyter manager"
        echo ""
        echo "  Commands:"
        echo "    status           Show all running Jupyter instances"
        echo "    stop-all         Stop all running instances"
        echo "    ports            Show port allocation (${PORT_RANGE_START}-${PORT_RANGE_END})"
        echo "    orphans          Find orphan Jupyter processes"
        echo "    kill-orphans     Kill orphan processes"
        echo "    which PORT       Show which project owns a port"
        echo "    config           Show current configuration"
        echo "    version          Show version"
        echo ""
        echo "  Configuration:"
        echo "    $JHUB_CONFIG_FILE"
        echo ""
        ;;
    *)
        echo "Unknown command: $1"
        echo "Run 'jupyter-hub help' for usage."
        exit 1
        ;;
esac