#!/usr/bin/env bash
set -euo pipefail

# deploy-docker.sh — Deploy Stonks Oracle to a Docker host via SSH
#
# Usage: bash deploy-docker.sh [OPTIONS]
#
# Options:
#   --host USER@HOST       SSH target (default: celes@192.168.42.254)
#   --ollama-url URL       Ollama API URL (default: auto-detect or install)
#   --ollama-model MODEL   Ollama model name (default: qwen3.5:9b-fast)
#   --dir PATH             Remote install directory (default: ~/stonks-oracle)
#
# Examples:
#   bash deploy-docker.sh
#   bash deploy-docker.sh --ollama-url http://10.1.1.12:2701 --ollama-model qwen3.6
#   bash deploy-docker.sh --host user@myserver --dir /opt/stonks

# -------------------------------------------------------
# Configuration (override via flags or environment)
# -------------------------------------------------------
REMOTE_HOST="${DEPLOY_HOST:-celes@192.168.42.254}"
REMOTE_DIR="${DEPLOY_DIR:-/home/celes/stonks-oracle}"
OLLAMA_URL="${DEPLOY_OLLAMA_URL:-}"
OLLAMA_MODEL="${DEPLOY_OLLAMA_MODEL:-qwen3.5:9b-fast}"
REPO_URL="http://admin:St0nks0racl3!@10.1.1.12:30300/admin/stonks-oracle.git"

# Parse command-line flags
while [[ $# -gt 0 ]]; do
  case $1 in
    --host) REMOTE_HOST="$2"; shift 2 ;;
    --ollama-url) OLLAMA_URL="$2"; shift 2 ;;
    --ollama-model) OLLAMA_MODEL="$2"; shift 2 ;;
    --dir) REMOTE_DIR="$2"; shift 2 ;;
    *) echo "Unknown option: $1"; exit 1 ;;
  esac
done

echo "=== Stonks Oracle Docker Deployment ==="
echo "  Target: ${REMOTE_HOST}:${REMOTE_DIR}"
echo "  Model:  ${OLLAMA_MODEL}"
echo "  Ollama: Docker container (GPU-accelerated)"
echo ""

# -------------------------------------------------------
# Step 0: Ensure prerequisites (multi-distro support)
# -------------------------------------------------------
echo "--- Step 0: Checking prerequisites ---"
ssh "$REMOTE_HOST" bash -s <<'REMOTE_SCRIPT'
set -euo pipefail

# --- Detect OS and package manager ---
detect_os() {
  if [ -f /etc/os-release ]; then
    . /etc/os-release
    OS_ID="${ID:-unknown}"
    OS_LIKE="${ID_LIKE:-$OS_ID}"
  elif [ -f /etc/redhat-release ]; then
    OS_ID="rhel"
    OS_LIKE="rhel"
  else
    OS_ID="unknown"
    OS_LIKE="unknown"
  fi

  # Detect WSL
  IS_WSL=false
  if grep -qi microsoft /proc/version 2>/dev/null; then
    IS_WSL=true
  fi

  # Determine package manager
  if command -v apt-get &>/dev/null; then
    PKG_MGR="apt"
  elif command -v dnf &>/dev/null; then
    PKG_MGR="dnf"
  elif command -v yum &>/dev/null; then
    PKG_MGR="yum"
  elif command -v pacman &>/dev/null; then
    PKG_MGR="pacman"
  elif command -v zypper &>/dev/null; then
    PKG_MGR="zypper"
  else
    PKG_MGR="unknown"
  fi

  echo "  Detected: OS=$OS_ID (like=$OS_LIKE), pkg=$PKG_MGR, WSL=$IS_WSL"
}

install_pkg() {
  local pkg="$1"
  case "$PKG_MGR" in
    apt)    sudo apt-get install -y "$pkg" ;;
    dnf)    sudo dnf -y install "$pkg" ;;
    yum)    sudo yum -y install "$pkg" ;;
    pacman) sudo pacman -S --noconfirm "$pkg" ;;
    zypper) sudo zypper install -y "$pkg" ;;
    *)      echo "  ERROR: Unknown package manager"; exit 1 ;;
  esac
}

update_pkg_cache() {
  case "$PKG_MGR" in
    apt)    sudo apt-get update -qq ;;
    dnf|yum) ;; # dnf/yum auto-refresh
    pacman) sudo pacman -Sy ;;
    zypper) sudo zypper refresh -q ;;
  esac
}

detect_os

# --- Git ---
if ! command -v git &>/dev/null; then
  echo "  Installing git..."
  update_pkg_cache
  install_pkg git
  echo "  ✓ Git installed"
else
  echo "  ✓ Git present"
fi

# --- Docker Engine ---
if command -v docker &>/dev/null && docker info &>/dev/null; then
  echo "  ✓ Docker already installed ($(docker --version | cut -d' ' -f3 | tr -d ','))"
else
  echo "  Installing Docker CE..."
  case "$PKG_MGR" in
    apt)
      # Debian/Ubuntu/WSL
      sudo apt-get update -qq
      sudo apt-get install -y ca-certificates curl gnupg
      sudo install -m 0755 -d /etc/apt/keyrings
      curl -fsSL https://download.docker.com/linux/${OS_ID}/gpg | sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg 2>/dev/null
      sudo chmod a+r /etc/apt/keyrings/docker.gpg
      echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/${OS_ID} $(. /etc/os-release && echo "$VERSION_CODENAME") stable" | \
        sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
      sudo apt-get update -qq
      sudo apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
      ;;
    dnf|yum)
      # RHEL/Rocky/Fedora/CentOS
      sudo "$PKG_MGR" -y install dnf-plugins-core 2>/dev/null || true
      local repo_distro="rhel"
      if [[ "$OS_ID" == "fedora" ]]; then repo_distro="fedora"; fi
      sudo dnf config-manager --add-repo "https://download.docker.com/linux/${repo_distro}/docker-ce.repo" 2>/dev/null || \
        sudo yum-config-manager --add-repo "https://download.docker.com/linux/${repo_distro}/docker-ce.repo" 2>/dev/null
      sudo "$PKG_MGR" -y install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
      ;;
    pacman)
      # Arch Linux
      sudo pacman -S --noconfirm docker docker-compose docker-buildx
      ;;
    zypper)
      # openSUSE
      sudo zypper install -y docker docker-compose docker-buildx
      ;;
  esac
  sudo systemctl enable --now docker 2>/dev/null || true
  sudo usermod -aG docker "$(whoami)" 2>/dev/null || true
  echo "  ✓ Docker installed and started"
fi

# --- Docker Compose plugin ---
if docker compose version &>/dev/null; then
  echo "  ✓ Docker Compose plugin available ($(docker compose version --short))"
else
  echo "  ERROR: docker compose plugin not found after Docker install"
  exit 1
fi

# --- NVIDIA Driver (skip on WSL — uses host driver) ---
if [ "$IS_WSL" = "true" ]; then
  echo "  ✓ WSL detected — using host Windows NVIDIA driver"
elif ! command -v nvidia-smi &>/dev/null; then
  echo "  Installing NVIDIA drivers..."
  case "$PKG_MGR" in
    apt)
      sudo apt-get install -y nvidia-driver-560 2>/dev/null || \
        sudo apt-get install -y nvidia-driver 2>/dev/null || \
        echo "  ⚠ NVIDIA driver install failed — install manually"
      ;;
    dnf|yum)
      sudo dnf -y install epel-release 2>/dev/null || true
      sudo dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel9/x86_64/cuda-rhel9.repo 2>/dev/null || true
      sudo dnf -y module install nvidia-driver:latest-dkms 2>/dev/null || \
        echo "  ⚠ NVIDIA driver install failed — install manually"
      ;;
    pacman)
      sudo pacman -S --noconfirm nvidia nvidia-utils 2>/dev/null || \
        echo "  ⚠ NVIDIA driver install failed — install manually"
      ;;
    zypper)
      echo "  ⚠ NVIDIA driver: install manually for openSUSE"
      ;;
  esac
else
  echo "  ✓ NVIDIA driver present ($(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -1))"
fi

# --- NVIDIA Container Toolkit ---
if command -v nvidia-ctk &>/dev/null; then
  echo "  ✓ NVIDIA Container Toolkit already installed"
elif [ "$IS_WSL" = "true" ] && docker run --rm --gpus all nvidia/cuda:12.8.0-base-ubuntu24.04 nvidia-smi &>/dev/null 2>&1; then
  echo "  ✓ WSL GPU passthrough working (no nvidia-ctk needed)"
else
  echo "  Installing NVIDIA Container Toolkit..."
  case "$PKG_MGR" in
    apt)
      curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg 2>/dev/null
      curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
        sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
        sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list > /dev/null
      sudo apt-get update -qq
      sudo apt-get install -y nvidia-container-toolkit
      ;;
    dnf|yum)
      curl -fsSL https://nvidia.github.io/libnvidia-container/stable/rpm/nvidia-container-toolkit.repo | \
        sudo tee /etc/yum.repos.d/nvidia-container-toolkit.repo > /dev/null
      sudo "$PKG_MGR" -y install nvidia-container-toolkit
      ;;
    pacman)
      sudo pacman -S --noconfirm nvidia-container-toolkit 2>/dev/null || \
        echo "  ⚠ Install nvidia-container-toolkit from AUR"
      ;;
    zypper)
      echo "  ⚠ NVIDIA Container Toolkit: install manually for openSUSE"
      ;;
  esac
  sudo nvidia-ctk runtime configure --runtime=docker 2>/dev/null || true
  sudo systemctl restart docker 2>/dev/null || true
  echo "  ✓ NVIDIA Container Toolkit installed and Docker configured"
fi

# --- Verify GPU is accessible from Docker ---
if docker run --rm --gpus all nvidia/cuda:12.8.0-base-ubuntu24.04 nvidia-smi &>/dev/null 2>&1; then
  echo "  ✓ GPU passthrough verified"
else
  echo "  ⚠ GPU passthrough test failed — may need a reboot or manual NVIDIA setup"
fi

# --- Firewall (open required ports if firewall is active) ---
if command -v firewall-cmd &>/dev/null && systemctl is-active firewalld &>/dev/null; then
  echo "  Configuring firewalld..."
  for port in 3000 8001 8002 8003 8004 9000 9001 11434; do
    sudo firewall-cmd --permanent --add-port="${port}/tcp" 2>/dev/null || true
  done
  sudo firewall-cmd --reload 2>/dev/null || true
  echo "  ✓ Firewall ports opened"
elif command -v ufw &>/dev/null && sudo ufw status 2>/dev/null | grep -q "active"; then
  echo "  Configuring ufw..."
  for port in 3000 8001 8002 8003 8004 9000 9001 11434; do
    sudo ufw allow "${port}/tcp" 2>/dev/null || true
  done
  echo "  ✓ UFW ports opened"
fi
REMOTE_SCRIPT
echo ""

# -------------------------------------------------------
# Step 1: Clone or update the repo on the remote host
# -------------------------------------------------------
echo "--- Step 1: Syncing repository ---"
ssh "$REMOTE_HOST" bash -s -- "$REMOTE_DIR" "$REPO_URL" <<'REMOTE_SCRIPT'
set -euo pipefail
REMOTE_DIR="$1"
REPO_URL="$2"

if [ -d "$REMOTE_DIR/.git" ]; then
  echo "  Updating existing repo..."
  cd "$REMOTE_DIR"
  git fetch origin
  git reset --hard origin/main
else
  echo "  Cloning fresh..."
  git clone "$REPO_URL" "$REMOTE_DIR"
  cd "$REMOTE_DIR"
fi

echo "  ✓ Repo synced at $(git log --oneline -1)"
REMOTE_SCRIPT
echo ""

# -------------------------------------------------------
# Step 2: Detect or configure Ollama
# -------------------------------------------------------
echo "--- Step 2: Configuring Ollama ---"
# Always use the Docker Ollama container with GPU passthrough
# The ollama/ollama image ships with CUDA runtime built-in
USE_DOCKER_OLLAMA=true
OLLAMA_URL="http://ollama:11434"
echo "  Using Docker Ollama container (GPU-accelerated via NVIDIA passthrough)"
echo "  Host-accessible at localhost:11434"
echo ""

# -------------------------------------------------------
# Step 3: Create .env and compose override
# -------------------------------------------------------
echo "--- Step 3: Configuring environment ---"
ssh "$REMOTE_HOST" bash -s -- "$REMOTE_DIR" "$OLLAMA_URL" "$OLLAMA_MODEL" "$USE_DOCKER_OLLAMA" <<'REMOTE_SCRIPT'
set -euo pipefail
REMOTE_DIR="$1"
OLLAMA_URL="$2"
OLLAMA_MODEL="$3"
USE_DOCKER_OLLAMA="$4"
cd "$REMOTE_DIR"

# Read API keys from local files if they exist
POLYGON_KEY=""
ALPACA_KEY=""
ALPACA_SECRET=""
ALPACA_URL="https://paper-api.alpaca.markets"

[ -f polygon.io.key ] && POLYGON_KEY=$(cat polygon.io.key)
[ -f alpaca.key ] && ALPACA_KEY=$(cat alpaca.key)
[ -f alpaca.secret ] && ALPACA_SECRET=$(cat alpaca.secret)
[ -f alpaca.url ] && ALPACA_URL=$(cat alpaca.url)

cat > .env <<EOF
# Stonks Oracle — Docker Deployment Environment
MARKET_DATA_API_KEY=${POLYGON_KEY}
BROKER_API_KEY=${ALPACA_KEY}
BROKER_API_SECRET=${ALPACA_SECRET}
BROKER_BASE_URL=${ALPACA_URL}
TRADING_ENABLED=true
TRADING_RISK_TIER=moderate
TRADING_MAX_OPEN_POSITIONS=15
OLLAMA_MODEL=${OLLAMA_MODEL}
MACRO_ENABLED=true
COMPETITIVE_ENABLED=true
EOF

# Create compose override based on Ollama configuration
if [ "$USE_DOCKER_OLLAMA" = "true" ]; then
  # Using Docker Ollama — no override needed, default compose handles it
  rm -f docker-compose.override.yml
  echo "  ✓ Using Docker Ollama container"
else
  # Using external Ollama — disable the container and point services to it
  # Determine if URL is localhost (needs host-gateway) or remote
  if echo "$OLLAMA_URL" | grep -qE "localhost|127\.0\.0\.1"; then
    DOCKER_OLLAMA_URL="http://host.docker.internal:$(echo "$OLLAMA_URL" | grep -oP ':\K[0-9]+')"
    cat > docker-compose.override.yml <<EOF
services:
  ollama:
    entrypoint: ["true"]
    restart: "no"
    ports: []
  extractor:
    depends_on:
      postgres:
        condition: service_healthy
      redis:
        condition: service_healthy
    environment:
      OLLAMA_BASE_URL: "${DOCKER_OLLAMA_URL}"
    extra_hosts:
      - "host.docker.internal:host-gateway"
  recommendation:
    environment:
      OLLAMA_BASE_URL: "${DOCKER_OLLAMA_URL}"
    extra_hosts:
      - "host.docker.internal:host-gateway"
EOF
  else
    # Remote Ollama — containers can reach it directly
    cat > docker-compose.override.yml <<EOF
services:
  ollama:
    entrypoint: ["true"]
    restart: "no"
  extractor:
    depends_on:
      postgres:
        condition: service_healthy
      redis:
        condition: service_healthy
    environment:
      OLLAMA_BASE_URL: "${OLLAMA_URL}"
  recommendation:
    environment:
      OLLAMA_BASE_URL: "${OLLAMA_URL}"
EOF
  fi
  echo "  ✓ Override created — services will use external Ollama at ${OLLAMA_URL}"
fi

echo "  ✓ .env configured (polygon=$([ -n "$POLYGON_KEY" ] && echo 'set' || echo 'empty'), alpaca=$([ -n "$ALPACA_KEY" ] && echo 'set' || echo 'empty'))"
REMOTE_SCRIPT
echo ""

# -------------------------------------------------------
# Step 4: Build and start all services
# -------------------------------------------------------
echo "--- Step 4: Building and starting services ---"
ssh "$REMOTE_HOST" bash -s -- "$REMOTE_DIR" "$USE_DOCKER_OLLAMA" <<'REMOTE_SCRIPT'
set -euo pipefail
REMOTE_DIR="$1"
USE_DOCKER_OLLAMA="$2"
cd "$REMOTE_DIR"

# Stop any existing deployment
docker compose down 2>/dev/null || true

# Build all images
echo "  Building images (this may take a few minutes)..."
docker compose build --quiet 2>&1 | tail -5

# Start infrastructure
echo "  Starting infrastructure..."
if [ "$USE_DOCKER_OLLAMA" = "true" ]; then
  docker compose up -d postgres redis minio minio-init ollama
else
  docker compose up -d postgres redis minio minio-init
fi

# Wait for infrastructure to be healthy
echo "  Waiting for infrastructure health checks..."
for svc in postgres redis minio; do
  for i in $(seq 1 30); do
    if docker compose ps "$svc" 2>/dev/null | grep -q healthy; then
      break
    fi
    sleep 2
  done
done
echo "  ✓ Infrastructure healthy"

# Start all application services
echo "  Starting application services..."
docker compose up -d

echo "  Waiting for services to stabilize..."
sleep 20

# Show status
echo ""
echo "  Service Status:"
docker compose ps --format "table {{.Name}}\t{{.Status}}" 2>/dev/null | head -25 || docker compose ps
REMOTE_SCRIPT
echo ""

# -------------------------------------------------------
# Step 5: Seed the database
# -------------------------------------------------------
echo "--- Step 5: Seeding database ---"
ssh "$REMOTE_HOST" bash -s -- "$REMOTE_DIR" <<'REMOTE_SCRIPT'
set -euo pipefail
cd "$1"

# Wait for query-api to be healthy
for i in $(seq 1 30); do
  if docker compose ps query-api 2>/dev/null | grep -q healthy; then
    break
  fi
  sleep 3
done

# Run the symbol registry seed
echo "  Seeding symbol registry..."
docker compose exec -T scheduler python -m services.symbol_registry.seed 2>/dev/null && echo "  ✓ Database seeded" || echo "  ⚠ Seed skipped (may already be seeded or service not ready)"
REMOTE_SCRIPT
echo ""

# -------------------------------------------------------
# Step 6: Ensure Ollama model is available
# -------------------------------------------------------
echo "--- Step 6: Checking Ollama model ---"
ssh "$REMOTE_HOST" bash -s -- "$OLLAMA_URL" "$OLLAMA_MODEL" "$USE_DOCKER_OLLAMA" "$REMOTE_DIR" <<'REMOTE_SCRIPT'
set -euo pipefail
OLLAMA_URL="$1"
OLLAMA_MODEL="$2"
USE_DOCKER_OLLAMA="$3"
REMOTE_DIR="$4"

if [ "$USE_DOCKER_OLLAMA" = "true" ]; then
  # Pull via Docker container
  cd "$REMOTE_DIR"
  if docker compose exec -T ollama ollama list 2>/dev/null | grep -q "$OLLAMA_MODEL"; then
    echo "  ✓ Model $OLLAMA_MODEL already available"
  else
    echo "  Pulling $OLLAMA_MODEL via Docker Ollama..."
    docker compose exec -T ollama ollama pull "$OLLAMA_MODEL"
    echo "  ✓ Model pulled"
  fi
else
  # Check via API
  if curl -sf "$OLLAMA_URL/api/tags" 2>/dev/null | grep -q "$OLLAMA_MODEL"; then
    echo "  ✓ Model $OLLAMA_MODEL already available at $OLLAMA_URL"
  else
    echo "  Pulling $OLLAMA_MODEL via $OLLAMA_URL..."
    curl -sf "$OLLAMA_URL/api/pull" -d "{\"name\":\"$OLLAMA_MODEL\"}" | tail -1
    echo "  ✓ Model pulled"
  fi
fi
REMOTE_SCRIPT
echo ""

# -------------------------------------------------------
# Done
# -------------------------------------------------------
REMOTE_IP=$(echo "$REMOTE_HOST" | cut -d@ -f2)
echo "=== Deployment Complete ==="
echo ""
echo "Endpoints:"
echo "  Dashboard:        http://${REMOTE_IP}:3000"
echo "  Query API:        http://${REMOTE_IP}:8004"
echo "  Symbol Registry:  http://${REMOTE_IP}:8001"
echo "  Trading Engine:   http://${REMOTE_IP}:8002"
echo "  Risk Engine:      http://${REMOTE_IP}:8003"
echo "  MinIO Console:    http://${REMOTE_IP}:9001"
echo "  Superset:         http://${REMOTE_IP}:8088"
echo "  Ollama:           http://${REMOTE_IP}:11434"
echo ""
echo "Commands:"
echo "  ssh $REMOTE_HOST 'cd $REMOTE_DIR && docker compose logs -f'"
echo "  ssh $REMOTE_HOST 'cd $REMOTE_DIR && docker compose ps'"
echo "  ssh $REMOTE_HOST 'cd $REMOTE_DIR && docker compose down'"