#!/usr/bin/env bash
set -u
# ============================================================
#  Streamlit LLM 应用 一键脚本（macOS / Linux / WSL）
#  公开源码，欢迎审查 —— 不放心可先复制给 AI 判断
#  默认配置：云间 API 中转站（https://cloudzone-api.cyou/）
#  生成：.streamlit/secrets.toml + chat_app.py（ChatGPT 式聊天应用）
# ============================================================

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"; }

CLOUDZONE_URL="https://cloudzone-api.cyou"
CLOUDZONE_API_BASE="https://cloudzone-api.cyou/v1"

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
}

check_python() {
  step "检查 Python 3..."
  if command -v python3 >/dev/null 2>&1 && python3 --version >/dev/null 2>&1; then
    PY="python3"
    info "找到 python3：$(python3 --version 2>&1)"
    return 0
  fi
  error "未找到 python3。请先安装：https://www.python.org/downloads/ （Windows 建议使用 WSL 后执行本脚本）"
  return 1
}

setup_env() {
  local net="$1"
  step "创建虚拟环境并安装 streamlit + openai（约 1-2 分钟）..."
  mkdir -p streamlit-app && cd streamlit-app || return 1
  python3 -m venv .venv 2>/dev/null || { error "创建虚拟环境失败，手动执行：python3 -m venv .venv"; return 1; }
  # shellcheck disable=SC1091
  source .venv/bin/activate || { warn "激活虚拟环境失败，继续使用系统环境"; }

  step "安装依赖..."
  local pypi="https://pypi.org/simple"
  if [ "$net" = "domestic" ]; then
    pypi="https://pypi.tuna.tsinghua.edu.cn/simple"
  fi

  if pip install streamlit openai -i "$pypi" --timeout 90 >/dev/null 2>&1; then
    info "依赖安装成功"
    return 0
  fi
  warn "主源失败，回退备用镜像（阿里云 PyPI）..."
  if pip install streamlit openai -i "https://mirrors.aliyun.com/pypi/simple" --timeout 90 >/dev/null 2>&1; then
    info "备用镜像安装成功"
    return 0
  fi
  error "依赖安装失败。请手动执行：pip install streamlit openai -i https://pypi.tuna.tsinghua.edu.cn/simple"
  return 1
}

configure_api() {
  step "配置 API Key..."
  echo ""
  echo "聊天应用需要一个 OpenAI 兼容的 Key。默认推荐云间 API 中转站："
  echo "  · OpenAI 兼容格式，代码无需改动，模型随便换"
  echo "  · 国内直连免翻墙，一把 Key 调 GPT / Claude / DeepSeek / GLM 等 90+ 模型"
  echo ""
  echo "是否现在跳转注册并获取 API Key？"
  echo "  Y - 立即跳转至 ${CLOUDZONE_URL}"
  echo "  N - 我已有 Key（OpenAI / DeepSeek 官方或第三方），自行输入"
  read -r -p "请选择 [Y/N]: " choice

  local API_KEY=""
  local BASE_URL="$CLOUDZONE_API_BASE"
  if [ "${choice:-N}" = "Y" ] || [ "${choice:-N}" = "y" ]; then
    info "正在打开浏览器..."
    if command -v xdg-open >/dev/null 2>&1; then
      xdg-open "$CLOUDZONE_URL" >/dev/null 2>&1
    elif command -v open >/dev/null 2>&1; then
      open "$CLOUDZONE_URL" >/dev/null 2>&1
    else
      warn "无法自动打开浏览器，请手动访问：$CLOUDZONE_URL"
    fi
    echo "注册后在控制台复制 Key（以 sk- 开头）。"
    read -r -p "粘贴你的 API Key: " API_KEY
    read -r -p "API 地址(直接回车用默认 ${CLOUDZONE_API_BASE}): " input_base
    [ -n "${input_base:-}" ] && BASE_URL="$input_base"
  else
    read -r -s -p "请输入你的 API Key (sk-...): " API_KEY
    echo
    echo "OpenAI 官方 Key 则地址填 https://api.openai.com/v1；DeepSeek 官方填 https://api.deepseek.com。"
    read -r -p "API 地址(直接回车用默认 ${CLOUDZONE_API_BASE}): " input_base
    [ -n "${input_base:-}" ] && BASE_URL="$input_base"
  fi

  if [ -z "${API_KEY:-}" ]; then
    error "未输入 API Key，跳过配置。可稍后手动编辑 .streamlit/secrets.toml。"
    return 1
  fi

  mkdir -p .streamlit
  cat > .streamlit/secrets.toml <<TOMEOF
# Streamlit secrets —— 请勿提交到 GitHub（含密钥）
OPENAI_API_KEY = "$API_KEY"
OPENAI_BASE_URL = "$BASE_URL"
TOMEOF

  info "已写入 .streamlit/secrets.toml（密钥文件，已自动加入 .gitignore 说明）"
  warn "注意：该文件含密钥，别提交到公开仓库！"

  if [ -f ".gitignore" ] && ! grep -q "secrets.toml" .gitignore; then
    echo ".streamlit/secrets.toml" >> .gitignore
  elif [ ! -f ".gitignore" ]; then
    echo ".streamlit/secrets.toml" > .gitignore
  fi
  return 0
}

write_app() {
  step "生成 chat_app.py（ChatGPT 式聊天应用）..."
  cat > chat_app.py <<'PYEOF'
import os
from openai import OpenAI
import streamlit as st

st.title("🤖 我的 AI 聊天室")

client = OpenAI(
    api_key=st.secrets.get("OPENAI_API_KEY", os.environ.get("OPENAI_API_KEY")),
    base_url=st.secrets.get("OPENAI_BASE_URL", os.environ.get("OPENAI_BASE_URL")),
)

if "model" not in st.session_state:
    st.session_state.model = "gpt-5.5"   # 换成支持模型名，如 deepseek-v4-pro / claude-sonnet-5

if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

if prompt := st.chat_input("说点什么…"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        stream = client.chat.completions.create(
            model=st.session_state.model,
            messages=[
                {"role": m["role"], "content": m["content"]}
                for m in st.session_state.messages
            ],
            stream=True,
        )
        response = st.write_stream(stream)
    st.session_state.messages.append({"role": "assistant", "content": response})
PYEOF
  info "已生成 streamlit-app/chat_app.py"
}

main() {
  echo "============================================"
  echo "  Streamlit LLM 应用 一键脚本"
  echo "  适用于 macOS / Linux / WSL"
  echo "  默认接入：云间 API 中转站"
  echo "============================================"
  echo ""

  local NET
  NET=$(detect_network)

  if ! check_python; then
    error "缺少 Python 3，脚本终止。请先安装 Python 后重跑。"
    exit 1
  fi

  if ! setup_env "$NET"; then
    error "环境准备失败，脚本终止。请按上方提示手动处理。"
    exit 1
  fi

  configure_api || warn "Key 未配置，应用仍会生成，运行时需自行补全 secrets"

  write_app

  echo ""
  echo "============================================"
  info "全部完成！启动你的 AI 聊天网页："
  echo "    cd streamlit-app && source .venv/bin/activate && streamlit run chat_app.py"
  echo "浏览器访问 http://localhost:8501"
  echo "换模型：改 chat_app.py 里 model 那一行"
  echo "============================================"
}

main "$@"