#!/bin/bash
# witty-compat — One-Click Installer for OpenCode
set -e

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'

CHECK="✓"
CROSS="✗"
STEP="→"

print_header() {
  echo -e "${BLUE}${BOLD}====================================================${NC}"
  echo -e "${CYAN}${BOLD}            witty-compat — Installer               ${NC}"
  echo -e "${BLUE}${BOLD}====================================================${NC}"
  echo ""
}

print_step() { echo -e "${STEP} ${BOLD}[$1/$2] $3...${NC}"; }
print_ok() { echo -e "${GREEN}  ${CHECK} $1${NC}"; }
print_err() { echo -e "${RED}  ${CROSS} $1${NC}"; }
print_warn() { echo -e "${YELLOW}  ! $1${NC}"; }

TOTAL_STEPS=9

# --- 1. Target detection ---
print_header
print_step 1 $TOTAL_STEPS "Detecting OpenCode"

if ! command -v opencode &> /dev/null; then
  print_err "opencode not found in PATH — install from https://opencode.ai/"
  exit 1
fi
print_ok "OpenCode detected"

# --- 2. Python & dependencies ---
print_step 2 $TOTAL_STEPS "Checking Python & dependencies"

if ! command -v python3 &> /dev/null; then
  print_err "python3 not found"
  exit 1
fi
PY_VER=$(python3 -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')")
print_ok "Python ${PY_VER} detected"

print_warn "Installing Python dependencies ..."
pip install pyyaml aiohttp jieba 2>&1 | while read line; do :; done
if python3 -c "import yaml, aiohttp, jieba" 2>/dev/null; then
  print_ok "pyyaml, aiohttp, jieba installed"
else
  print_err "Python dependency install failed — try: pip install pyyaml aiohttp jieba"
fi

# --- 3. Determine paths ---
print_step 3 $TOTAL_STEPS "Preparing directories"

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CONFIG_DIR="${OPENCODE_CONFIG_DIR:-${HOME}/.config/opencode}"
PLUGIN_DIR="${CONFIG_DIR}/plugins"
SKILLS_DIR="${CONFIG_DIR}/skills"
WITTY_COMPAT_DIR="${CONFIG_DIR}/witty-compat"

mkdir -p "${PLUGIN_DIR}" "${SKILLS_DIR}"

# If running from a different location, copy there. Otherwise link in-place.
if [ "${SCRIPT_DIR}" != "${WITTY_COMPAT_DIR}" ]; then
  if [ -e "${WITTY_COMPAT_DIR}" ]; then
    if [ -L "${WITTY_COMPAT_DIR}" ]; then
      rm -f "${WITTY_COMPAT_DIR}"
    else
      print_warn "${WITTY_COMPAT_DIR} exists, backing up to ${WITTY_COMPAT_DIR}.bak"
      mv "${WITTY_COMPAT_DIR}" "${WITTY_COMPAT_DIR}.bak"
    fi
  fi
  ln -s "${SCRIPT_DIR}" "${WITTY_COMPAT_DIR}"
  print_ok "Symlinked witty-compat → ${WITTY_COMPAT_DIR}"
else
  print_ok "Using in-place installation: ${WITTY_COMPAT_DIR}"
fi

# --- 4. SearXNG MCP (auto-start Docker) ---
print_step 4 $TOTAL_STEPS "Search MCP (searxng)"

SEARXNG_PORT=8888
SETTINGS_YML="${WITTY_COMPAT_DIR}/searxng/settings.yml"

if command -v docker &> /dev/null; then
  docker rm -f witty-searxng 2>/dev/null || true
  for port in $(seq 8888 8987); do
    if ss -tlnp 2>/dev/null | grep -q ":${port} "; then
      continue
    fi
    if curl -s --max-time 2 "http://127.0.0.1:${port}/search?q=test&format=json" > /dev/null 2>&1; then
      print_ok "SearXNG already running at port ${port}"
      SEARXNG_PORT=$port
      break
    fi
    if docker run -d --name witty-searxng \
      --dns 223.5.5.5 --dns 114.114.114.114 \
      -p "127.0.0.1:${port}:8080" \
      -v "${SETTINGS_YML}:/etc/searxng/settings.yml:ro" \
      --restart unless-stopped \
      searxng/searxng > /dev/null 2>&1; then
      echo -n "  waiting for SearXNG on port ${port}"
      for i in $(seq 1 20); do
        sleep 1; echo -n "."
        if curl -s --max-time 3 "http://127.0.0.1:${port}/search?q=test&format=json" > /dev/null 2>&1; then
          echo ""
          print_ok "SearXNG started on port ${port}"
          SEARXNG_PORT=$port
          break 2
        fi
      done
      echo ""
      print_warn "SearXNG did not respond on port ${port}"
    fi
  done
  if [ "$SEARXNG_PORT" = "8888" ] && [ "$port" != "8888" ]; then
    print_warn "SearXNG not started — ports 8888-8987 unavailable"
  fi
else
  print_warn "docker not found — skipping SearXNG"
  echo -e "  Install Docker for web search, or set SEARXNG_URL to an existing instance."
fi

# --- 5. Register plugin ---
print_step 5 $TOTAL_STEPS "Registering plugin"

PLUGIN_LINK="${PLUGIN_DIR}/witty-compat.js"
PLUGIN_SRC="${WITTY_COMPAT_DIR}/.opencode/plugins/witty-compat.js"

if [ ! -f "${PLUGIN_SRC}" ]; then
  print_err "Plugin source not found: ${PLUGIN_SRC}"
  exit 1
fi

rm -f "${PLUGIN_LINK}"
ln -s "${PLUGIN_SRC}" "${PLUGIN_LINK}"
print_ok "Plugin registered → ${PLUGIN_LINK}"

# --- 6. Register agents ---
print_step 6 $TOTAL_STEPS "Registering agents"

OPENCODE_CONFIG="${CONFIG_DIR}/opencode.json"
AGENT_CONFIG="${WITTY_COMPAT_DIR}/.opencode/witty-compat.json"

if [ ! -f "${OPENCODE_CONFIG}" ]; then
  print_warn "opencode.json not found at ${OPENCODE_CONFIG}, skipping agent registration"
else
  python3 - "$OPENCODE_CONFIG" "$AGENT_CONFIG" "$WITTY_COMPAT_DIR" "$SEARXNG_PORT" <<'PYEOF'
import json, sys

base_path = sys.argv[1]
agent_path = sys.argv[2]
install_dir = sys.argv[3]
searxng_port = sys.argv[4]

with open(base_path) as f:
    base = json.load(f)
with open(agent_path) as f:
    agent_cfg = f.read()

agent_cfg = agent_cfg.replace('__WITTY_COMPAT_HOME__', install_dir)
agent_cfg = agent_cfg.replace('"http://127.0.0.1:8888"', f'"http://127.0.0.1:{searxng_port}"')
agent_cfg_data = json.loads(agent_cfg)

base.setdefault("agent", {}).update(agent_cfg_data.get("agent", {}))
for mcpid, mcpdef in agent_cfg_data.get("mcp", {}).items():
    if mcpid not in base.get("mcp", {}):
        base.setdefault("mcp", {})[mcpid] = mcpdef

with open(base_path, "w") as f:
    json.dump(base, f, indent=2, ensure_ascii=False)
PYEOF
  print_ok "Agents registered in opencode.json"
fi

# --- 7. Symlink skills ---
print_step 7 $TOTAL_STEPS "Symlinking skills"

SKILLS_LINK="${SKILLS_DIR}/witty-compat"
SKILLS_SRC="${WITTY_COMPAT_DIR}/skills/witty-compat"

rm -rf "${SKILLS_LINK}"
ln -s "${SKILLS_SRC}" "${SKILLS_LINK}"
print_ok "Skills symlinked → ${SKILLS_LINK}"

# --- 8. Bootstrap knowledge base ---
print_step 8 $TOTAL_STEPS "Bootstrapping knowledge base"

CASES_DIR="${WITTY_CASES_DIR:-${WITTY_COMPAT_DIR}/knowledge/cases}"
CASES_REPO="https://atomgit.com/openeuler/witty-ops-cases.git"

if [ -d "${CASES_DIR}" ] && [ -n "$(ls -A "${CASES_DIR}" 2>/dev/null)" ]; then
  print_ok "Case library present: ${CASES_DIR}"
elif command -v git &> /dev/null; then
  print_warn "Cloning case library (${CASES_REPO}) ..."
  if git clone --depth 1 "${CASES_REPO}" "${CASES_DIR}" > /dev/null 2>&1; then
    print_ok "Case library cloned → ${CASES_DIR}"
  else
    print_warn "Clone failed (offline?). Skipping — only built-in branches will be indexed."
  fi
else
  print_warn "git not found, skipping case library. Only built-in branches will be indexed."
fi

if python3 "${WITTY_COMPAT_DIR}/tools/kb_search.py" --build > /dev/null 2>&1; then
  print_ok "Search index built"
else
  print_warn "Index build failed — will build lazily on first search"
fi

# --- 9. Verify tools ---
print_step 9 $TOTAL_STEPS "Verifying tools & dependencies"

VERIFY_OK=true

python3 "${WITTY_COMPAT_DIR}/tools/benchmark.py" --help > /dev/null 2>&1 && \
  print_ok "benchmark.py OK" || { print_err "benchmark.py check failed"; VERIFY_OK=false; }

python3 "${WITTY_COMPAT_DIR}/tuner/orchestrator.py" --list-modules > /dev/null 2>&1 && \
  print_ok "orchestrator.py OK" || { print_err "orchestrator.py check failed"; VERIFY_OK=false; }

# System tools (best-effort, not required)
echo ""
echo -e "${CYAN}─────────────────────────────────────────────────────${NC}"
echo -e "${CYAN}  System tools${NC}"
echo -e "${CYAN}─────────────────────────────────────────────────────${NC}"
for tool in curl numactl nvidia-smi npu-smi; do
  if command -v $tool &> /dev/null; then
    print_ok "$tool detected"
  else
    print_warn "$tool not found — install if needed:"
    case $tool in
      curl)     echo -e "      yum install curl  /  apt install curl" ;;
      numactl)  echo -e "      yum install numactl  /  apt install numactl" ;;
      nvidia-smi) echo -e "      需安装 NVIDIA 驱动 (nvidia-driver)" ;;
      npu-smi)  echo -e "      需安装 Ascend CANN (toolkit)" ;;
    esac
  fi
done

# --- Done ---
echo ""
echo -e "${BLUE}─────────────────────────────────────────────────────${NC}"
if [ "$VERIFY_OK" = true ]; then
  echo -e "${GREEN}${BOLD}  All done! witty-compat installed successfully.${NC}"
  echo ""
  echo -e "  Restart OpenCode and select the ${CYAN}compat-Agent${NC} to get started."
  echo -e "  Or ask: ${CYAN}\"帮我部署 vLLM 推理服务\"${NC}"
else
  echo -e "${YELLOW}${BOLD}  Installation complete but some tool checks failed.${NC}"
  echo -e "  Review errors above and ensure Python dependencies are installed."
fi
echo ""
echo -e "${BLUE}====================================================${NC}"