#!/usr/bin/env bash
set -u
# ============================================================
#  Ollama 本地部署大模型 一键脚本（macOS / Linux / WSL）
#  公开源码，欢迎审查 —— 不放心可先复制给 AI 判断
#  依赖：本脚本仅安装 Ollama 本身，模型需另行 ollama run
# ============================================================

GREEN='\033[0;32m'; YELLOW='\033[1;33m'; RED='\033[0;31m'; CYAN='\033[0;36m'; NC='\033[0m'
info()  { echo -e "${GREEN}[INFO]${NC} $1"; }
warn()  { echo -e "${YELLOW}[WARN]${NC} $1"; }
error() { echo -e "${RED}[ERROR]${NC} $1"; }
step()  { echo -e "${CYAN}[STEP]${NC} $1"; }

detect_network() {
  step "检测网络环境（国内 / 国外）..."
  if curl -fsI --max-time 5 "https://claude.ai" >/dev/null 2>&1; then
    info "可直连 claude.ai，判定为海外网络"
    echo "overseas"
  else
    warn "无法直连 claude.ai，判定为国内网络环境（将自动切换国内镜像）"
    echo "domestic"
  fi
}

install_official() {
  # 官方安装脚本
  step "使用官方安装脚本安装 Ollama..."
  if curl -fsSL --max-time 60 https://ollama.com/install.sh | sh; then
    info "官方脚本安装成功"
    return 0
  fi
  warn "官方脚本安装失败，回退备用方案（手动下载二进制）..."
  return 1
}

install_binary() {
  local net="$1"
  step "备用方案：从 GitHub Releases 下载 Ollama 二进制..."
  local sys="$(uname -s)"
  local asset=""
  case "$sys" in
    Darwin) asset="Ollama-darwin.zip" ;;
    Linux)  asset="ollama-linux-amd64.tgz" ;;
    *)
      error "未知系统 $sys，请手动到 https://ollama.com/download 下载安装包"
      return 1
      ;;
  esac

  local base_url="https://github.com/ollama/ollama/releases/latest/download/$asset"
  if [ "$net" = "domestic" ]; then
    base_url="https://ghfast.top/https://github.com/ollama/ollama/releases/latest/download/$asset"
  fi

  local tmp="/tmp/ollama-dl"
  mkdir -p "$tmp"
  if curl -fsSL --max-time 300 "$base_url" -o "$tmp/$asset"; then
    info "下载成功（$asset）"
  else
    warn "主源失败，回退备用源（ghfast.top 镜像）..."
    base_url="https://ghfast.top/https://github.com/ollama/ollama/releases/latest/download/$asset"
    if ! curl -fsSL --max-time 300 "$base_url" -o "$tmp/$asset"; then
      error "二进制下载失败。请手动访问 https://ollama.com/download 下载安装"
      return 1
    fi
  fi

  local dest="$HOME/ollama"
  mkdir -p "$dest"
  case "$asset" in
    *.zip)
      if command -v unzip >/dev/null 2>&1; then
        unzip -qo "$tmp/$asset" -d "$dest" || { error "解压失败"; return 1; }
      else
        ditto -x -k "$tmp/$asset" "$dest" 2>/dev/null || { error "解压失败（需要 unzip 或 ditto）"; return 1; }
      fi
      ;;
    *.tgz)
      tar -xzf "$tmp/$asset" -C "$dest" || { error "解压失败"; return 1; }
      ;;
  esac

  # 加入 PATH（当前会话立即生效，永久写入 ~/.bashrc / ~/.zshrc；原因：ollama 需要能找到）
  local bin_dir="$dest/bin"
  export PATH="$PATH:$bin_dir"
  local rc="$HOME/.bashrc"
  [ "$SHELL" = "/bin/zsh" ] && rc="$HOME/.zshrc"
  if grep -qF "$bin_dir" "$rc" 2>/dev/null; then
    :
  else
    echo "" >> "$rc"
    echo "# Ollama - 由一键脚本写入" >> "$rc"
    echo "export PATH=\"\$PATH:$bin_dir\"" >> "$rc"
  fi
  info "已安装到 $dest，执行 ollama 前请先运行：source $rc（或新开终端）"
  return 0
}

main() {
  echo "============================================"
  echo "  Ollama 本地部署大模型 一键脚本"
  echo "  适用于 macOS / Linux / WSL"
  echo "============================================"
  echo ""

  local NET
  NET=$(detect_network)

  if command -v ollama >/dev/null 2>&1; then
    info "检测到 Ollama 已安装：$(ollama --version 2>&1)"
  else
    if ! install_official; then
      if ! install_binary "$NET"; then
        error "安装失败，脚本终止。请手动到 https://ollama.com/download 安装后重跑。"
        exit 1
      fi
    fi

    if ! command -v ollama >/dev/null 2>&1; then
      # 备用二进制安装时可能未在 PATH；定位提示
      if [ -x "$HOME/ollama/bin/ollama" ]; then
        export PATH="$PATH:$HOME/ollama/bin"
      fi
    fi
  fi

  if command -v ollama >/dev/null 2>&1; then
    info "Ollama 安装完成：$(ollama --version 2>&1)"
  else
    error "未能在当前会话找到 ollama 命令。请 source ~/.bashrc（或 ~/.zshrc）后重试，或手动安装。"
    exit 1
  fi

  echo ""
  echo "============================================"
  info "全部完成！现在运行你的第一个本地模型："
  echo "    ollama run gemma4"
  echo "（首次运行自动下载模型，约数 GB，请耐心等待）"
  echo "退出聊天：输入 /bye"
  echo "============================================"
}

main "$@"