#!/bin/bash
set -euo pipefail

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

info()    { echo -e "${CYAN}[INFO]${NC} $*"; }
success() { echo -e "${GREEN}[OK]${NC}   $*"; }
warn()    { echo -e "${YELLOW}[WARN]${NC} $*"; }
error()   { echo -e "${RED}[ERR]${NC}  $*" >&2; }
step()    { echo -e "\n${BOLD}${BLUE}══ $* ${NC}"; }

# ─── OS Detection ─────────────────────────────────────────────────────────────
detect_os() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        OS_ID="${ID:-unknown}"
        OS_ID_LIKE="${ID_LIKE:-}"
    else
        error "Cannot detect OS (no /etc/os-release found)"
        exit 1
    fi

    case "$OS_ID" in
        aosc)   DISTRO="aosc" ;;
        debian) DISTRO="debian" ;;
        ubuntu) DISTRO="ubuntu" ;;
        fedora) DISTRO="fedora" ;;
        *)
            # fallback via ID_LIKE
            case "$OS_ID_LIKE" in
                *debian*) DISTRO="debian" ;;
                *fedora*|*rhel*) DISTRO="fedora" ;;
                *) error "Unsupported distro: $OS_ID"; exit 1 ;;
            esac
            ;;
    esac
    info "Detected distro: ${BOLD}$DISTRO${NC} (ID=$OS_ID)"
}

# ─── HTTP Proxy ───────────────────────────────────────────────────────────────
setup_proxy() {
    step "HTTP Proxy"
    echo -e "Do you want to configure an HTTP proxy for this session? ${YELLOW}(helps with Homebrew downloads)${NC}"
    read -rp "Configure proxy? [y/N] " ans
    case "$ans" in
        [Yy]*)
            while true; do
                read -rp "Enter proxy URL (e.g. http://192.168.1.1:7890): " proxy_url
                if [[ "$proxy_url" =~ ^https?://[^:]+:[0-9]+$ ]]; then
                    break
                fi
                warn "Invalid format. Expected: http://<IP>:<PORT> or https://<IP>:<PORT>"
            done
            export http_proxy="$proxy_url"
            export https_proxy="$proxy_url"
            export HTTP_PROXY="$proxy_url"
            export HTTPS_PROXY="$proxy_url"
            success "Proxy set to $proxy_url for this session"
            ;;
        *)
            info "Skipping proxy configuration"
            ;;
    esac
}

# ─── SSH Key Setup ────────────────────────────────────────────────────────────
setup_ssh_key() {
    step "SSH Key Configuration"

    # Create ~/.ssh
    mkdir -p ~/.ssh
    chmod 700 ~/.ssh
    touch ~/.ssh/authorized_keys
    chmod 600 ~/.ssh/authorized_keys

    # Show already-present keys
    existing_count=$(grep -c . ~/.ssh/authorized_keys 2>/dev/null || true)
    if [ "$existing_count" -gt 0 ]; then
        info "Existing authorized keys ($existing_count):"
        grep -v '^\s*$' ~/.ssh/authorized_keys | while IFS= read -r line; do
            echo "    ${line:0:60}..."
        done
    fi

    # Collect public key(s) interactively
    echo ""
    echo -e "Paste ${BOLD}new${NC} SSH public key(s) to add (enter blank line to finish):"
    echo ""

    local added=0
    while true; do
        read -rp "Public key (or blank to finish): " pubkey
        [[ -z "$pubkey" ]] && break
        if [[ "$pubkey" =~ ^(ssh-rsa|ssh-ed25519|ecdsa-sha2-nistp256|sk-ssh-ed25519) ]]; then
            if grep -qF "$pubkey" ~/.ssh/authorized_keys 2>/dev/null; then
                info "Key already present, skipping (${pubkey:0:30}...)"
            else
                echo "$pubkey" >> ~/.ssh/authorized_keys
                success "Key added (${pubkey:0:30}...)"
                (( added++ )) || true
            fi
        else
            warn "Does not look like a valid public key, skipping"
        fi
    done

    if [ "$added" -eq 0 ]; then
        info "No new keys added"
    else
        success "authorized_keys updated (+$added key(s))"
    fi

    # Configure sshd_config
    info "Configuring /etc/ssh/sshd_config ..."
    SSHD_CONF="/etc/ssh/sshd_config"

    sudo_set_sshd() {
        local key="$1" val="$2"
        # Uncomment or add the line
        if sudo grep -qE "^\s*#?\s*${key}\s" "$SSHD_CONF"; then
            sudo sed -i -E "s|^\s*#?\s*(${key})\s+.*|\1 ${val}|" "$SSHD_CONF"
        else
            echo "${key} ${val}" | sudo tee -a "$SSHD_CONF" > /dev/null
        fi
    }

    sudo_set_sshd "PubkeyAuthentication"  "yes"
    sudo_set_sshd "AuthorizedKeysFile"    ".ssh/authorized_keys"

    # Only disable password auth when at least one key is present — avoid lockout
    if grep -qv '^\s*$' ~/.ssh/authorized_keys 2>/dev/null; then
        sudo_set_sshd "PasswordAuthentication" "no"
        success "sshd_config updated (password login disabled)"
    else
        warn "No authorized keys found — keeping PasswordAuthentication unchanged to avoid lockout"
    fi

    # Restart SSH
    if sudo systemctl restart ssh 2>/dev/null || sudo systemctl restart sshd 2>/dev/null; then
        success "SSH service restarted"
    else
        warn "Could not restart SSH service automatically — please restart it manually"
    fi
}

# ─── Install git via system package manager ───────────────────────────────────
install_git() {
    if command -v git &>/dev/null; then
        success "git already installed ($(git --version))"
        return
    fi
    info "Installing git via system package manager ..."
    case "$DISTRO" in
        aosc)           sudo oma install -y git ;;
        debian|ubuntu)  sudo apt-get update -qq && sudo apt-get install -y git ;;
        fedora)         sudo dnf install -y git ;;
    esac
    success "git installed"
}

# ─── Homebrew ─────────────────────────────────────────────────────────────────
install_homebrew() {
    if command -v brew &>/dev/null; then
        success "Homebrew already installed"
        return
    fi
    info "Installing Homebrew ..."
    local install_script
    install_script="$(mktemp)"
    if ! curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh -o "$install_script"; then
        rm -f "$install_script"
        error "Failed to download Homebrew install script (network error)"
        exit 1
    fi
    /bin/bash "$install_script"
    rm -f "$install_script"

    # Add brew to PATH for the rest of this script
    if [ -f /home/linuxbrew/.linuxbrew/bin/brew ]; then
        eval "$(/home/linuxbrew/.linuxbrew/bin/brew shellenv)"
    fi

    if ! command -v brew &>/dev/null; then
        error "Homebrew installation failed (brew not found after install)"
        exit 1
    fi
    success "Homebrew installed"
}

# brew install，跳过已安装的包，忽略依赖 post-install 失败引起的非零退出，最后校验每个包
brew_install_packages() {
    local to_install=()
    for pkg in "$@"; do
        if brew list --formula "$pkg" &>/dev/null; then
            info "Homebrew: $pkg already installed, skipping"
        else
            to_install+=("$pkg")
        fi
    done

    if [ ${#to_install[@]} -eq 0 ]; then
        success "All Homebrew packages already installed"
        return
    fi

    info "Installing via Homebrew: ${to_install[*]}"
    brew install "${to_install[@]}" || true

    local failed=()
    for pkg in "${to_install[@]}"; do
        if ! brew list --formula "$pkg" &>/dev/null; then
            failed+=("$pkg")
        fi
    done
    if [ ${#failed[@]} -gt 0 ]; then
        error "The following packages failed to install: ${failed[*]}"
        exit 1
    fi
    success "Homebrew packages installed: ${to_install[*]}"
}

# ─── Software Installation ────────────────────────────────────────────────────
install_packages() {
    step "Software Installation"
    case "$DISTRO" in
        aosc)
            info "Installing packages via oma ..."
            sudo oma install -y git fish eza fastfetch btop docker docker-compose docker-buildx
            success "All packages installed via oma"
            ;;
        debian|ubuntu)
            install_git
            install_homebrew
            brew_install_packages fish eza fastfetch btop
            ;;
        fedora)
            install_git
            install_homebrew
            brew_install_packages fish eza fastfetch btop
            ;;
    esac
}

# ─── Fish Shell Setup ─────────────────────────────────────────────────────────
setup_fish() {
    step "Fish Shell Configuration"

    FISH_PATH="$(command -v fish)"
    if [ -z "$FISH_PATH" ]; then
        error "fish not found in PATH"
        return 1
    fi

    # Add fish to /etc/shells if not already present
    if ! grep -qF "$FISH_PATH" /etc/shells; then
        echo "$FISH_PATH" | sudo tee -a /etc/shells > /dev/null
        success "Added $FISH_PATH to /etc/shells"
    else
        info "$FISH_PATH already in /etc/shells"
    fi

    # Change default shell only if not already fish
    current_shell="$(getent passwd "$USER" | cut -d: -f7)"
    if [ "$current_shell" = "$FISH_PATH" ]; then
        info "fish is already the default shell"
    else
        sudo chsh -s "$FISH_PATH" "$USER"
        success "Default shell changed to fish ($FISH_PATH)"
    fi

    # For brew-based systems: add brew shellenv to fish config
    if [ "$DISTRO" != "aosc" ] && [ -f /home/linuxbrew/.linuxbrew/bin/brew ]; then
        mkdir -p ~/.config/fish
        FISH_CONF="$HOME/.config/fish/config.fish"
        BREW_LINE='eval (/home/linuxbrew/.linuxbrew/bin/brew shellenv)'
        if ! grep -qF "$BREW_LINE" "$FISH_CONF" 2>/dev/null; then
            echo "$BREW_LINE" >> "$FISH_CONF"
            success "brew shellenv added to fish config"
        else
            info "brew shellenv already in fish config"
        fi
    fi
}

# ─── Docker Installation ──────────────────────────────────────────────────────
install_docker() {
    step "Docker Installation"

    case "$DISTRO" in
        aosc)
            info "Docker already installed via oma, skipping"
            ;;
        debian|ubuntu)
            if command -v docker &>/dev/null; then
                info "Docker already installed ($(docker --version)), skipping"
            else
                curl -fsSL https://git.mitsea.com/FlintyLemming/scripts-public/raw/branch/main/linux-managements/install-docker.sh \
                    -o /tmp/install-docker.sh
                sudo sh /tmp/install-docker.sh
                success "Docker installed"
            fi
            ;;
        fedora)
            if command -v docker &>/dev/null; then
                info "Docker already installed ($(docker --version)), skipping"
            else
                info "Setting up Docker CE repository ..."
                sudo curl -fsSL https://download.docker.com/linux/fedora/docker-ce.repo \
                    -o /etc/yum.repos.d/docker-ce.repo
                sudo dnf install -y docker-ce docker-ce-cli containerd.io \
                    docker-compose-plugin docker-buildx-plugin
                success "Docker installed"
            fi
            ;;
    esac

    docker_no_root
}

# Allow current user to run docker without sudo
docker_no_root() {
    info "Configuring Docker for non-root usage ..."

    if ! getent group docker > /dev/null 2>&1; then
        sudo groupadd docker
    fi

    if id -nG "$USER" | grep -qw docker; then
        info "User '$USER' is already in the docker group"
    else
        sudo usermod -aG docker "$USER"
        success "User '$USER' added to the docker group"
        warn "Log out and back in for the group change to take effect"
    fi

    if ! sudo systemctl is-enabled --quiet docker 2>/dev/null; then
        sudo systemctl enable docker
    fi
    sudo systemctl start docker
    success "Docker service running"
}

# ─── Dotfiles ─────────────────────────────────────────────────────────────────
clone_dotfiles() {
    step "Dotfiles"
    DOTFILES_DIR="$HOME/.flinty"
    if [ -d "$DOTFILES_DIR" ]; then
        info "Dotfiles directory already exists, pulling latest ..."
        git -C "$DOTFILES_DIR" pull
    else
        info "Cloning dotfiles ..."
        git clone https://git.mitsea.com/FlintyLemming/dotfiles.git "$DOTFILES_DIR"
    fi
    success "Dotfiles ready at $DOTFILES_DIR"
}

# ─── Fish Config (source dotfiles) ───────────────────────────────────────────
configure_fish_dotfiles() {
    step "Fish Config — Dotfiles Integration"
    mkdir -p ~/.config/fish
    FISH_CONF="$HOME/.config/fish/config.fish"
    SOURCE_LINE="source ~/.flinty/fish/add-on.fish"
    if ! grep -qF "$SOURCE_LINE" "$FISH_CONF" 2>/dev/null; then
        echo "$SOURCE_LINE" >> "$FISH_CONF"
        success "Added dotfiles source to fish config"
    else
        info "Dotfiles source line already in fish config"
    fi
}

# ─── SSH Config ───────────────────────────────────────────────────────────────
configure_ssh_config() {
    step "SSH Config — Dotfiles Integration"
    mkdir -p ~/.ssh
    SSH_CONF="$HOME/.ssh/config"
    INCLUDE_LINE="Include ~/.flinty/ssh/config"
    if ! grep -qF "$INCLUDE_LINE" "$SSH_CONF" 2>/dev/null; then
        # Include must be at the top of ssh config
        if [ -f "$SSH_CONF" ]; then
            tmp="$(mktemp)"
            { echo "$INCLUDE_LINE"; echo ""; cat "$SSH_CONF"; } > "$tmp"
            mv "$tmp" "$SSH_CONF"
        else
            echo "$INCLUDE_LINE" > "$SSH_CONF"
        fi
        chmod 600 "$SSH_CONF"
        success "Include line added to ~/.ssh/config"
    else
        info "Include line already in ~/.ssh/config"
    fi
}

# ─── Sudo Privilege Check ─────────────────────────────────────────────────────
ensure_sudo() {
    step "Sudo Privilege Check"

    if sudo -v 2>/dev/null; then
        success "User '$USER' has sudo access"
        # Keep sudo ticket alive in the background for the duration of the script
        ( while true; do sudo -n true 2>/dev/null; sleep 50; done ) &
        SUDO_KEEPALIVE_PID=$!
        return
    fi

    warn "User '$USER' does not have sudo access"

    case "$DISTRO" in
        aosc|fedora) local sudo_group="wheel" ;;
        debian|ubuntu) local sudo_group="sudo" ;;
    esac

    info "Enter the ${BOLD}root password${NC} to grant sudo access (adds '$USER' to '$sudo_group' group):"
    if su -c "usermod -aG $sudo_group $USER" root; then
        success "User '$USER' added to '$sudo_group' group"
        echo ""
        warn "Group change requires re-login to take effect."
        warn "Please log out, log back in, and re-run this script."
        exit 0
    else
        error "Failed to grant sudo. Please run as root: usermod -aG $sudo_group $USER"
        exit 1
    fi
}

# ─── Main ─────────────────────────────────────────────────────────────────────
main() {
    echo -e "${BOLD}${CYAN}"
    echo "╔══════════════════════════════════════════╗"
    echo "║         FlintyLemming Server Setup       ║"
    echo "╚══════════════════════════════════════════╝"
    echo -e "${NC}"

    detect_os
    ensure_sudo

    setup_proxy
    setup_ssh_key
    install_packages
    setup_fish
    install_docker
    clone_dotfiles
    configure_fish_dotfiles
    configure_ssh_config

    step "Starting Docker"
    if sudo systemctl start docker; then
        success "Docker started"
    else
        warn "Could not start Docker — please start it manually"
    fi

    # Stop sudo keepalive
    if [ -n "${SUDO_KEEPALIVE_PID:-}" ]; then
        kill "$SUDO_KEEPALIVE_PID" 2>/dev/null || true
    fi

    echo ""
    echo -e "${BOLD}${GREEN}══ Setup complete! ══${NC}"
    echo -e "Please ${BOLD}log out and back in${NC} (or open a new shell) for all changes to take effect."
}

main "$@"