#!/bin/bash
# Cloudflare employee laptop setup for OpenCode
# Inspired by gitlab.cfdata.org/srhea/agents-for-all
#
# Installs prerequisites, developer tools, OpenCode (CLI/TUI + Desktop),
# authenticates with the AI gateway, and installs a smart shell wrapper
# that handles re-authentication automatically.
#
# Smart wrapper based on gitlab.cfdata.org/-/snippets/799 by Trey Guinn

set -u

GREEN=$'\e[0;32m'
RED=$'\e[0;31m'
YELLOW=$'\e[93m'
BLUE=$'\e[34m'
NC=$'\e[0m'

OUTPUT=/dev/null

CLOUDFLARE_USER=${USER:-}

readonly SSH_DIR="${HOME}/.ssh"
readonly CLOUDFLARE_SSH_DIR="${SSH_DIR}/cloudflare"
readonly KEY_TYPE="ed25519"
readonly ZSHRC="${HOME}/.zshrc"

SSH_PUB_KEY="${CLOUDFLARE_SSH_DIR}/id_${KEY_TYPE}.pub"
SSH_PRIVATE_KEY="${CLOUDFLARE_SSH_DIR}/id_${KEY_TYPE}"

mkdir -p "$SSH_DIR"
chmod 0711 "$SSH_DIR"
mkdir -p "$CLOUDFLARE_SSH_DIR"
chmod 0711 "$CLOUDFLARE_SSH_DIR"

green_echo() { echo "${GREEN}$1${NC}"; }
red_echo() { echo "${RED}$1${NC}"; }
yellow_echo() { echo "${YELLOW}$1${NC}"; }
blue_echo() { echo "${BLUE}$1${NC}"; }

ensure_shell_config() {
    if [[ ! -f "$ZSHRC" ]]; then
        touch "$ZSHRC"
        green_echo "   Created .zshrc"
    fi
}

welcome() {
    clear
    echo "========================================"
    green_echo "   Cloudflare OpenCode Setup"
    echo "========================================"
    echo
    echo "This script prepares your laptop for OpenCode."
    echo "It installs developer tools, both OpenCode"
    echo "interfaces (CLI/TUI and Desktop), and a smart"
    echo "wrapper that handles re-authentication."
    echo
}

set_cloudflare_user() {
    CLOUDFLARE_USER=${CLOUDFLARE_USER/@cloudflare.com/}
    green_echo "   User: ${CLOUDFLARE_USER}"
    echo
}

install_brew() {
    echo "Checking for Homebrew..."

    local brew_path="/usr/local/bin/brew"
    if [[ $(uname -m) == "arm64" ]]; then
        brew_path="/opt/homebrew/bin/brew"
    fi

    if ! command -v brew >$OUTPUT; then
        yellow_echo "   Installing Homebrew..."
        /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
    fi

    if ! grep -q "shellenv" "$ZSHRC"; then
        echo "eval \"\$($brew_path shellenv)\"" >> "$ZSHRC"
    fi

    eval "$($brew_path shellenv)"

    green_echo "   Homebrew ready"

    if ! command -v jq >$OUTPUT; then
        brew install jq >$OUTPUT
    fi
}

install_cloudflared() {
    echo "Installing cloudflared..."

    if [[ $(uname -m) == "arm64" ]]; then
        if ! (arch -arch x86_64 uname -m >$OUTPUT); then
            yellow_echo "   Installing Rosetta..."
            /usr/sbin/softwareupdate --install-rosetta --agree-to-license
        fi
    fi

    if ! command -v cloudflared >$OUTPUT; then
        brew install cloudflared >$OUTPUT
    fi

    green_echo "   cloudflared ready"
}

generate_ssh_key() {
    echo "Checking SSH keys..."

    if [ -f "${SSH_PRIVATE_KEY}" ]; then
        if [ -f "${SSH_PUB_KEY}" ]; then
            green_echo "   SSH key pair exists"
            chmod 600 "${SSH_PRIVATE_KEY}"
            ssh-add --apple-use-keychain "${SSH_PRIVATE_KEY}" 2>/dev/null
            return
        else
            ssh-keygen -y -f "${SSH_PRIVATE_KEY}" >"${SSH_PUB_KEY}"
            chmod 600 "${SSH_PRIVATE_KEY}"
            ssh-add --apple-use-keychain "${SSH_PRIVATE_KEY}" 2>/dev/null
            green_echo "   Public key regenerated"
            return
        fi
    fi

    if [ -f "${SSH_DIR}/id_ed25519" ] && [ -f "${SSH_DIR}/id_ed25519.pub" ]; then
        green_echo "   Found existing SSH key: ${SSH_DIR}/id_ed25519"
        read -p "   Use this key for Cloudflare? [Y/n] " -n 1 -r REPLY
        echo
        if [[ $REPLY =~ ^[Yy]$ ]] || [[ -z $REPLY ]]; then
            cp "${SSH_DIR}/id_ed25519" "${SSH_PRIVATE_KEY}"
            cp "${SSH_DIR}/id_ed25519.pub" "${SSH_PUB_KEY}"
            chmod 600 "${SSH_PRIVATE_KEY}"
            ssh-add --apple-use-keychain "${SSH_PRIVATE_KEY}" 2>/dev/null
            green_echo "   Existing key configured"
            return
        fi
    fi

    read -p "   Generate a new SSH key? [Y/n] " -n 1 -r REPLY
    echo
    if [[ $REPLY =~ ^[Yy]$ ]] || [[ -z $REPLY ]]; then
        ssh-keygen -t "$KEY_TYPE" -C "${CLOUDFLARE_USER}@cloudflare.com" -f "${SSH_PRIVATE_KEY}" -N "" >/dev/null
        green_echo "   Key generated"
        ssh-add --apple-use-keychain "${SSH_PRIVATE_KEY}" 2>/dev/null
    else
        yellow_echo "   Skipped key generation"
    fi
}

gitlab_key_upload() {
    if [ ! -f "${SSH_PUB_KEY}" ]; then
        yellow_echo "   No public key found, skipping GitLab upload"
        return
    fi

    echo "========================================"
    echo "ACTION REQUIRED: Upload SSH key to GitLab"
    echo "========================================"
    echo
    yellow_echo "   Your public key (copied to clipboard):"
    echo "   ---"
    cat "${SSH_PUB_KEY}"
    echo "   ---"
    echo

    pbcopy <"${SSH_PUB_KEY}"
    green_echo "   Key copied to clipboard"

    echo "   Opening GitLab SSH settings..."
    sleep 2
    open "https://gitlab.cfdata.org/-/user_settings/ssh_keys"

    echo
    echo "   Paste the key (Cmd+V) into GitLab and click 'Add key'."
    echo
    read -n 1 -s -r -p "   Press any key once done..."
    echo
    green_echo "   Continuing..."
}

create_ssh_config() {
    echo "Writing SSH config..."
    local ssh_config="${SSH_DIR}/config"
    local cloudflare_config="${CLOUDFLARE_SSH_DIR}/config"
    local homebrew_prefix
    homebrew_prefix="$(brew --prefix)"

    cat >"$cloudflare_config" <<EOF
Host *.cfdata.org
    AddKeysToAgent yes
    UseKeychain yes
    User ${CLOUDFLARE_USER}
    IdentityFile ${SSH_PRIVATE_KEY}
    StrictHostKeyChecking no

Host *.ssh.cfdata.org
    ProxyCommand ${homebrew_prefix}/bin/cloudflared access ssh --hostname %h
EOF

    if [[ ! -f "$ssh_config" ]]; then touch "$ssh_config"; fi

    cp "${ssh_config}" "${ssh_config}.setup"
    sed '/CLOUDFLARE SETUP/d' <"${ssh_config}.setup" >"${ssh_config}"

    cat >>"$ssh_config" <<EOF

# CLOUDFLARE SETUP
Match all # CLOUDFLARE SETUP
    Include ${CLOUDFLARE_SSH_DIR}/config # CLOUDFLARE SETUP
EOF
    rm "${ssh_config}.setup"
    green_echo "   SSH config updated"
}

setup_git_identity() {
    echo "Checking git identity..."
    if git config --global user.email | grep -q @cloudflare.com$; then
        green_echo "   Git identity already set"
    else
        git config --global user.email "${CLOUDFLARE_USER}@cloudflare.com"
        git config --global user.name "${CLOUDFLARE_USER}"
        green_echo "   Git identity set to ${CLOUDFLARE_USER}@cloudflare.com"
    fi
}

install_internal_tools() {
    echo "Installing internal tools..."
    if brew tap cloudflare/engineering ssh://git@gitlab.cfdata.org/cloudflare/pkg/homebrew.git 2>/dev/null; then
        brew install cfsetup >$OUTPUT 2>&1 || true
        brew install cf-paste >$OUTPUT 2>&1 || true
        green_echo "   Internal tools installed"
    else
        red_echo "   Could not access internal tap (check SSH key upload)"
    fi
}

install_warp_certs() {
    echo "Installing WARP certificates..."
    if brew install cloudflare-certs >$OUTPUT 2>&1; then
        yellow_echo "   May prompt for your laptop password..."
        install-cloudflare-warp-certs
        green_echo "   Certificates installed"
    else
        yellow_echo "   Skipped (tap may be missing)"
    fi
}

install_opencode() {
    echo "========================================"
    echo "Installing OpenCode..."
    echo "========================================"

    if [[ $(uname -m) == "arm64" ]]; then eval "$(/opt/homebrew/bin/brew shellenv)"; fi

    # CLI/TUI: the 'opencode' terminal command
    if ! command -v opencode >$OUTPUT; then
        yellow_echo "   Installing OpenCode CLI/TUI..."
        brew install opencode >$OUTPUT
    else
        green_echo "   OpenCode CLI/TUI already installed"
    fi

    # Desktop: the native GUI app
    if [[ -d "/Applications/OpenCode.app" ]] || [[ -d "${HOME}/Applications/OpenCode.app" ]]; then
        green_echo "   OpenCode Desktop already installed"
    else
        yellow_echo "   Installing OpenCode Desktop..."
        brew install --cask opencode-desktop >$OUTPUT
    fi

    echo
    green_echo "   OpenCode CLI/TUI: run 'opencode' in any terminal"
    green_echo "   OpenCode Desktop: find it in Applications or Spotlight"
    echo "   Both share the same config (~/.config/opencode/opencode.jsonc)"
    echo
}

authenticate_ai_gateway() {
    echo "========================================"
    echo "Authenticating with Cloudflare AI gateway..."
    echo "========================================"
    echo
    yellow_echo "   A browser window will open for Cloudflare SSO."
    yellow_echo "   Complete the login with your YubiKey, then return here."
    echo

    if opencode auth login https://opencode.cloudflare.dev; then
        green_echo "   AI gateway authenticated"
    else
        red_echo "   Auth failed. You can retry later with:"
        echo "   opencode auth login https://opencode.cloudflare.dev"
    fi
}

install_smart_wrapper() {
    echo "Installing smart auth wrapper..."

    # Check if wrapper is already installed
    if grep -q "OPENCODE_SMART_WRAPPER" "$ZSHRC" 2>/dev/null; then
        green_echo "   Smart wrapper already installed"
        return
    fi

    cat >> "$ZSHRC" << 'WRAPPER_EOF'

# OPENCODE_SMART_WRAPPER_START
# Smart OpenCode launcher with automatic re-authentication
# Based on gitlab.cfdata.org/-/snippets/799 by Trey Guinn
# Installed by cf2.opencode.school setup script
opencode() {
    local RED='\033[0;31m'
    local GREEN='\033[0;32m'
    local YELLOW='\033[1;33m'
    local BLUE='\033[0;34m'
    local NC='\033[0m'

    local REQUIRED_MCPS=("cf-portal")
    local TOKEN_MAX_AGE=82800  # 23 hours in seconds

    log_info() { echo -e "${BLUE}[INFO]${NC} $1" }
    log_success() { echo -e "${GREEN}[SUCCESS]${NC} $1" }
    log_warn() { echo -e "${YELLOW}[WARN]${NC} $1" }
    log_error() { echo -e "${RED}[ERROR]${NC} $1" }

    get_token_file() {
        local pattern="$HOME/.cloudflared/opencode.cloudflare.dev-*-token"
        local files=($~pattern)
        if [[ ${#files[@]} -gt 0 && -f "${files[1]}" ]]; then
            echo "${files[1]}"
        fi
    }

    needs_auth() {
        local token_file=$(get_token_file)
        if [[ -z "$token_file" || ! -f "$token_file" ]]; then
            return 0
        fi
        local file_age
        if [[ "$(uname)" == "Darwin" ]]; then
            file_age=$(($(date +%s) - $(stat -f %m "$token_file")))
        else
            file_age=$(($(date +%s) - $(stat -c %Y "$token_file")))
        fi
        if [[ $file_age -ge $TOKEN_MAX_AGE ]]; then
            return 0
        fi
        return 1
    }

    auth_opencode() {
        log_info "Authenticating OpenCode with Cloudflare Access..."
        echo ""
        if command opencode auth login https://opencode.cloudflare.dev; then
            log_success "Successfully authenticated OpenCode"
            return 0
        else
            log_error "Failed to authenticate OpenCode"
            return 1
        fi
    }

    get_mcp_status() {
        local mcp_name=$1
        local mcp_list_output
        mcp_list_output=$(command opencode mcp list 2>&1 || true)
        if echo "$mcp_list_output" | grep -A1 "●.*$mcp_name " | grep -q "connected"; then
            echo "connected"
        elif echo "$mcp_list_output" | grep -A1 "●.*$mcp_name " | grep -q "failed"; then
            echo "failed"
        elif echo "$mcp_list_output" | grep -A1 "●.*$mcp_name " | grep -q "disabled"; then
            echo "disabled"
        else
            echo "unknown"
        fi
    }

    auth_mcp() {
        local mcp_name=$1
        log_info "Authenticating MCP: $mcp_name..."
        if command opencode mcp auth "$mcp_name"; then
            log_success "Successfully authenticated $mcp_name"
            return 0
        else
            log_error "Failed to authenticate $mcp_name"
            return 1
        fi
    }

    get_token_remaining_time() {
        local token_file=$(get_token_file)
        if [[ -z "$token_file" || ! -f "$token_file" ]]; then
            echo "unknown"
            return
        fi
        local file_age
        if [[ "$(uname)" == "Darwin" ]]; then
            file_age=$(($(date +%s) - $(stat -f %m "$token_file")))
        else
            file_age=$(($(date +%s) - $(stat -c %Y "$token_file")))
        fi
        local remaining=$((TOKEN_MAX_AGE - file_age))
        if [[ $remaining -lt 0 ]]; then
            echo "expired"
        else
            local hours=$((remaining / 3600))
            local minutes=$(((remaining % 3600) / 60))
            echo "${hours}h ${minutes}m"
        fi
    }

    log_info "Starting OpenCode launch sequence..."
    echo ""

    if needs_auth; then
        log_warn "Token expired or missing - authentication required"
        if ! auth_opencode; then
            log_error "Cannot proceed without OpenCode authentication"
            return 1
        fi
    else
        local remaining=$(get_token_remaining_time)
        log_success "Token still valid (${remaining} remaining) - skipping auth"
    fi

    echo ""
    log_info "Checking MCP authentication status..."
    echo ""

    local needs_mcp_auth=()
    local mcp mcp_status
    for mcp in "${REQUIRED_MCPS[@]}"; do
        mcp_status=$(get_mcp_status "$mcp")
        case "$mcp_status" in
            "connected")
                log_success "$mcp is already connected"
                ;;
            "failed"|"disabled")
                log_warn "$mcp is $mcp_status - will attempt to authenticate"
                needs_mcp_auth+=("$mcp")
                ;;
            "unknown")
                log_warn "$mcp status unknown - will attempt to authenticate"
                needs_mcp_auth+=("$mcp")
                ;;
        esac
    done

    if [[ ${#needs_mcp_auth[@]} -gt 0 ]]; then
        echo ""
        log_info "Authenticating MCPs that need setup..."
        echo ""
        for mcp in "${needs_mcp_auth[@]}"; do
            if ! auth_mcp "$mcp"; then
                log_warn "Continuing despite $mcp authentication failure..."
            fi
            echo ""
        done
    fi

    echo ""
    log_success "All checks complete!"
    echo ""

    log_info "Launching OpenCode..."
    command opencode "$@"
}
# OPENCODE_SMART_WRAPPER_END
WRAPPER_EOF

    green_echo "   Smart wrapper installed in ~/.zshrc"
    green_echo "   Every time you run 'opencode', it checks and refreshes auth automatically"
}

all_done() {
    echo "========================================"
    green_echo "   Setup complete!"
    echo "========================================"
    echo
    echo "Your laptop now has:"
    echo "   Homebrew"
    echo "   SSH access to GitLab"
    echo "   Cloudflare Tunnel (cloudflared)"
    echo "   Internal certificates"
    echo "   Internal tools (cfsetup, cf-paste)"
    echo "   OpenCode CLI/TUI"
    echo "   OpenCode Desktop"
    echo "   AI gateway authenticated"
    echo "   Smart auth wrapper (handles re-authentication)"
    echo
    echo "The smart wrapper runs every time you type 'opencode'"
    echo "in a terminal. It checks your auth tokens and refreshes"
    echo "them automatically if they have expired."
    echo
    echo "========================================"
    echo "Next: return to cf2.opencode.school"
    echo "to continue with the course."
    echo "========================================"
    echo
    yellow_echo "Restart your terminal or run 'source ~/.zshrc' to activate the wrapper."
}

main() {
    ensure_shell_config
    welcome
    set_cloudflare_user

    install_brew

    if ssh -T git@gitlab.cfdata.org >/dev/null 2>&1; then
        green_echo "   SSH to GitLab already works"
    else
        generate_ssh_key
        gitlab_key_upload
    fi

    create_ssh_config
    setup_git_identity
    install_cloudflared
    install_internal_tools
    install_warp_certs
    install_opencode
    authenticate_ai_gateway
    install_smart_wrapper
    all_done
}

main