#!/bin/bash
# ============================================================================
# DaSiWa API Server — Автоматическая установка
# Запуск: chmod +x setup.sh && sudo ./setup.sh
# ============================================================================

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
API_DIR="/root/custom_comfyui"
API_PORT=8080

echo ""
echo "============================================================"
echo "🚀 DaSiWa API Server — Установка"
echo "============================================================"
echo ""

# ============================================================================
# 1. Проверка root
# ============================================================================
if [ "$EUID" -ne 0 ]; then
    echo "❌ Запусти от root: sudo ./setup.sh"
    exit 1
fi

# ============================================================================
# 2. Обновление пакетов
# ============================================================================
echo "📦 Обновление пакетов..."
apt-get update -qq

# ============================================================================
# 3. Python 3 + pip
# ============================================================================
echo "🐍 Проверка Python..."
if command -v python3 &>/dev/null; then
    PYTHON_VERSION=$(python3 --version 2>&1)
    echo "   ✅ $PYTHON_VERSION"
else
    echo "   📥 Установка Python 3..."
    apt-get install -y python3 python3-pip python3-venv
fi

if ! command -v pip3 &>/dev/null; then
    apt-get install -y python3-pip
fi

# ============================================================================
# 4. wget
# ============================================================================
if ! command -v wget &>/dev/null; then
    echo "📥 Установка wget..."
    apt-get install -y wget
fi

# ============================================================================
# 5. Копирование файлов API
# ============================================================================
echo ""
echo "📂 Настройка API директории: $API_DIR"

if [ "$SCRIPT_DIR" != "$API_DIR" ]; then
    mkdir -p "$API_DIR"
    cp "$SCRIPT_DIR/server.py" "$API_DIR/"
    cp "$SCRIPT_DIR/hmac_auth.py" "$API_DIR/"
    cp "$SCRIPT_DIR/requirements.txt" "$API_DIR/"
    cp "$SCRIPT_DIR/generate_keys.py" "$API_DIR/"

    [ -f "$SCRIPT_DIR/keys.json" ] && cp "$SCRIPT_DIR/keys.json" "$API_DIR/" && echo "   ✅ keys.json скопирован"
    [ -f "$SCRIPT_DIR/workflow_api.json" ] && cp "$SCRIPT_DIR/workflow_api.json" "$API_DIR/" && echo "   ✅ workflow_api.json скопирован"

    echo "   ✅ Файлы скопированы в $API_DIR"
else
    echo "   ✅ Уже в правильной директории"
fi

# ============================================================================
# 6. Установка Python зависимостей
# ============================================================================
echo ""
echo "📦 Установка Python зависимостей..."
pip3 install -r "$API_DIR/requirements.txt" --quiet

# ============================================================================
# 7. Генерация ключей (если нет)
# ============================================================================
if [ ! -f "$API_DIR/keys.json" ]; then
    echo ""
    echo "🔐 Генерация ключей..."
    cd "$API_DIR"
    python3 generate_keys.py
    echo ""
    echo "⚠️  ВАЖНО: Скопируй keys.json на свой ПК (клиент)!"
    echo "   scp root@<ip_сервера>:$API_DIR/keys.json ."
else
    echo "🔐 keys.json найден — используем существующие ключи"
fi

# ============================================================================
# 8. Проверка workflow
# ============================================================================
if [ ! -f "$API_DIR/workflow_api.json" ]; then
    echo ""
    echo "⚠️  workflow_api.json не найден!"
    echo "   Загрузи его: scp workflow_api.json root@<ip>:$API_DIR/"
fi

# ============================================================================
# 9. Настройка Firewall (UFW)
# ============================================================================
echo ""
echo "🔥 Настройка Firewall..."
if ! command -v ufw &>/dev/null; then
    apt-get install -y ufw
fi

ufw allow 22/tcp     # SSH
ufw allow ${API_PORT}/tcp  # API
# НЕ открываем 8188 — ComfyUI только через localhost
ufw --force enable
echo "   ✅ UFW: открыты порты 22 (SSH), ${API_PORT} (API)"
echo "   🔒 Порт 8188 (ComfyUI) закрыт снаружи"

# ============================================================================
# 10. Установка systemd сервиса (автозапуск)
# ============================================================================
echo ""
echo "⚙️  Настройка systemd сервиса..."

cat > /etc/systemd/system/dasiwa-api.service << EOF
[Unit]
Description=DaSiWa API Server
After=network.target

[Service]
Type=simple
User=root
WorkingDirectory=$API_DIR
Environment=COMFY_HOST=127.0.0.1
Environment=COMFY_PORT=8188
Environment=API_PORT=${API_PORT}
Environment=COMFY_OUTPUT_DIR=/ComfyUI/output
ExecStart=/usr/bin/python3 $API_DIR/server.py
Restart=on-failure
RestartSec=5

[Install]
WantedBy=multi-user.target
EOF

systemctl daemon-reload
systemctl enable dasiwa-api
systemctl start dasiwa-api
echo "   ✅ Сервис dasiwa-api включён и запущен"
echo "   ✅ Автозапуск при старте системы: ON"

# ============================================================================
# ГОТОВО
# ============================================================================
echo ""
echo "============================================================"
echo "✅ УСТАНОВКА ЗАВЕРШЕНА"
echo "============================================================"
echo ""
echo "📋 Чеклист:"
echo "   1. Загрузи модели в /ComfyUI/models/ (если ещё нет)"
echo "   2. Загрузи workflow_api.json в $API_DIR/"
echo "   3. Скопируй keys.json на свой ПК"
echo ""
echo "🔧 Управление:"
echo "   systemctl status dasiwa-api    — статус"
echo "   systemctl restart dasiwa-api   — перезапуск"
echo "   journalctl -u dasiwa-api -f    — логи"
echo ""
echo "🧪 Проверка:"
echo "   curl http://localhost:${API_PORT}/health"
echo ""
echo "🔐 API доступен на:"
echo "   http://<ip_сервера>:${API_PORT}/run"
echo "============================================================"