#!/bin/bash
set -euo pipefail

# InfraPilot one-line installer.
#
#   curl -fsSL https://get.infrapilot.sh | bash                         # Community (free, keyless)
#   curl -fsSL https://get.infrapilot.sh | bash -s -- --license IP-...  # Enterprise (in place)
#
# Re-running updates in place. If a paid license is present (flag, env, or a prior
# install's .env) the script fetches private-registry pull creds from the license
# server, logs in, and runs the Enterprise image — reusing the SAME data volume, so
# CE→EE upgrades preserve your database and config. No key = Community, no signup.

INFRAPILOT_VERSION="${INFRAPILOT_VERSION:-latest}"
INSTALL_DIR="${INSTALL_DIR:-$HOME/infrapilot}"
LICENSE_KEY="${LICENSE_KEY:-}"
DO_UNINSTALL=0
PURGE=0
BASE_URL="${INFRAPILOT_BASE_URL:-https://infrapilot.org}"
CE_COMPOSE_URL="$BASE_URL/docker-compose.install.yml"
EE_COMPOSE_URL="$BASE_URL/docker-compose.install.ee.yml"
EE_TOKEN_URL="$BASE_URL/api/license/ee-pull-token"

RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; BLUE='\033[0;34m'; NC='\033[0m'

# ── Args ─────────────────────────────────────────────────────────────────────
while [ $# -gt 0 ]; do
    case "$1" in
        --license)   LICENSE_KEY="${2:-}"; shift 2 ;;
        --license=*) LICENSE_KEY="${1#*=}"; shift ;;
        --dir)       INSTALL_DIR="${2:-}"; shift 2 ;;
        --dir=*)     INSTALL_DIR="${1#*=}"; shift ;;
        --version)   INFRAPILOT_VERSION="${2:-}"; shift 2 ;;
        --version=*) INFRAPILOT_VERSION="${1#*=}"; shift ;;
        --uninstall) DO_UNINSTALL=1; shift ;;
        --purge)     PURGE=1; shift ;;
        -h|--help)
            echo "Usage: install.sh [--license KEY] [--dir PATH] [--version TAG]"
            echo "       install.sh --uninstall [--purge]   # remove InfraPilot (--purge also deletes the data volume)"
            exit 0 ;;
        *) echo -e "${RED}Unknown option: $1${NC}"; exit 1 ;;
    esac
done

# ── Uninstall ────────────────────────────────────────────────────────────────
if [ "$DO_UNINSTALL" = "1" ]; then
    command -v docker &> /dev/null || { echo -e "${RED}docker not found${NC}"; exit 1; }
    if [ ! -d "$INSTALL_DIR" ]; then
        echo -e "${YELLOW}Nothing to uninstall at ${INSTALL_DIR}.${NC}"; exit 0
    fi
    if [ "$PURGE" = "1" ]; then
        echo -e "${BLUE}→ Uninstalling InfraPilot and deleting its data volume...${NC}"
        docker compose --project-directory "$INSTALL_DIR" down -v 2> /dev/null || true
    else
        echo -e "${BLUE}→ Uninstalling InfraPilot (keeping the data volume)...${NC}"
        docker compose --project-directory "$INSTALL_DIR" down 2> /dev/null || true
    fi
    rm -rf "$INSTALL_DIR"
    echo -e "${GREEN}✓ InfraPilot uninstalled.${NC}"
    if [ "$PURGE" = "1" ]; then
        echo "  Data volume removed."
    else
        echo "  Data volume kept — re-run with ${YELLOW}--uninstall --purge${NC} to delete your data too."
    fi
    exit 0
fi

echo ""
echo -e "${BLUE}  ___        __          ____  _ _       _   ${NC}"
echo -e "${BLUE} |_ _|_ __  / _|_ __ __ |  _ \\(_) | ___ | |_ ${NC}"
echo -e "${BLUE}  | || '_ \\| |_| '__/ _\`| |_) | | |/ _ \\| __|${NC}"
echo -e "${BLUE}  | || | | |  _| | | (_| |  __/| | | (_) | |_ ${NC}"
echo -e "${BLUE} |___|_| |_|_| |_|  \\__,_|_|  |_|_|\\___/ \\__|${NC}"
echo ""

# ── Prereqs ──────────────────────────────────────────────────────────────────
check_command() {
    if ! command -v "$1" &> /dev/null; then
        echo -e "${RED}Error: $1 is required but not installed.${NC}"
        echo "  Install $1 and try again."
        exit 1
    fi
}
check_command docker
check_command curl
if ! docker compose version &> /dev/null; then
    echo -e "${RED}Error: Docker Compose v2 is required.${NC}"
    echo "  Install Docker Desktop or the Docker Compose plugin."
    exit 1
fi

# Extract a JSON string field without requiring jq.
json_get() {
    printf '%s' "$1" | grep -o "\"$2\"[[:space:]]*:[[:space:]]*\"[^\"]*\"" | head -1 \
        | sed "s/.*:[[:space:]]*\"\\(.*\\)\"/\\1/"
}
gen_secret() {
    if command -v openssl &> /dev/null; then openssl rand -hex 32
    else (set +o pipefail; LC_ALL=C tr -dc 'A-Za-z0-9' < /dev/urandom | head -c 64); fi
}
gen_uuid() {
    cat /proc/sys/kernel/random/uuid 2>/dev/null \
        || uuidgen 2>/dev/null \
        || openssl rand -hex 16 2>/dev/null \
        || echo "inst-$(date +%s)-$RANDOM"
}

# ── Carry over values from an existing install ───────────────────────────────
mkdir -p "$INSTALL_DIR"
cd "$INSTALL_DIR"

existing_key=""; existing_jwt=""; existing_iid=""
if [ -f .env ]; then
    existing_key=$(grep '^LICENSE_KEY=' .env | cut -d= -f2- || true)
    existing_jwt=$(grep '^JWT_SECRET=' .env | cut -d= -f2- || true)
    existing_iid=$(grep '^INSTANCE_ID=' .env | cut -d= -f2- || true)
fi
[ -z "$LICENSE_KEY" ] && [ -n "$existing_key" ] && LICENSE_KEY="$existing_key"
JWT_SECRET="${JWT_SECRET:-${existing_jwt:-$(gen_secret)}}"
INSTANCE_ID="${existing_iid:-$(gen_uuid)}"

# ── Resolve edition: paid license → Enterprise, else Community ───────────────
MODE="ce"
EE_IMAGE=""
if [ -n "$LICENSE_KEY" ]; then
    echo -e "${BLUE}→ Checking license...${NC}"
    hn="$(hostname 2>/dev/null || echo unknown)"
    resp=$(curl -fsS -m 20 -X POST "$EE_TOKEN_URL" \
        -H 'content-type: application/json' \
        -d "{\"key\":\"$LICENSE_KEY\",\"instance_id\":\"$INSTANCE_ID\",\"hostname\":\"$hn\",\"version\":\"$INFRAPILOT_VERSION\"}" \
        2>/dev/null || true)
    token=$(json_get "$resp" token)
    if [ -n "$token" ]; then
        reg=$(json_get "$resp" registry)
        user=$(json_get "$resp" username)
        img=$(json_get "$resp" image)
        ver=$(json_get "$resp" version)
        tier=$(json_get "$resp" tier)
        if echo "$token" | docker login "$reg" -u "$user" --password-stdin &> /dev/null; then
            MODE="ee"
            EE_IMAGE="$img"
            [ -n "$ver" ] && INFRAPILOT_VERSION="$ver"
            echo -e "${GREEN}  ✓ Enterprise Edition unlocked (${tier}).${NC}"
        else
            echo -e "${YELLOW}  ! Registry login failed — installing Community Edition.${NC}"
        fi
    else
        err=$(json_get "$resp" error)
        echo -e "${YELLOW}  ! ${err:-Could not reach the license server} — installing Community Edition.${NC}"
    fi
fi

if [ "$MODE" = "ce" ]; then
    echo -e "${GREEN}  Community Edition${NC}  (free · 1 server)"
fi
echo -e "  Version: ${INFRAPILOT_VERSION}"
echo ""

# ── Write config + compose, then start ───────────────────────────────────────
echo -e "${BLUE}→ Downloading configuration...${NC}"
if [ "$MODE" = "ee" ]; then
    curl -fsSL "$EE_COMPOSE_URL" -o docker-compose.yml
else
    curl -fsSL "$CE_COMPOSE_URL" -o docker-compose.yml
fi

{
    echo "JWT_SECRET=${JWT_SECRET}"
    echo "LICENSE_KEY=${LICENSE_KEY}"
    echo "INSTANCE_ID=${INSTANCE_ID}"
    echo "INFRAPILOT_VERSION=${INFRAPILOT_VERSION}"
    [ "$MODE" = "ee" ] && echo "EE_IMAGE=${EE_IMAGE}"
} > .env
chmod 600 .env

echo -e "${BLUE}→ Pulling InfraPilot image...${NC}"
docker compose pull

echo -e "${BLUE}→ Starting InfraPilot...${NC}"
docker compose up -d

echo -e "${BLUE}→ Waiting for InfraPilot to be ready...${NC}"
for _ in $(seq 1 30); do
    if curl -sf http://localhost/api/health > /dev/null 2>&1; then break; fi
    sleep 2
done

echo ""
echo -e "${GREEN}================================================${NC}"
if [ "$MODE" = "ee" ]; then
    echo -e "${GREEN}  InfraPilot Enterprise is running!${NC}"
else
    echo -e "${GREEN}  InfraPilot is running!${NC}"
fi
echo -e "${GREEN}================================================${NC}"
echo ""
echo "  Dashboard:  http://localhost"
echo "  Data dir:   $INSTALL_DIR"
echo ""
if [ "$MODE" = "ce" ]; then
    echo -e "  ${BLUE}Unlock more servers, scanning & SSO:${NC}"
    echo "    1. Buy a plan at $BASE_URL/pricing"
    echo "    2. Re-run:  curl -fsSL https://get.infrapilot.sh | bash -s -- --license YOUR_KEY"
    echo "    (upgrades in place — your data is preserved)"
    echo ""
fi
echo "  To update:  curl -fsSL https://get.infrapilot.sh | bash"
echo "  To stop:      docker compose --project-directory $INSTALL_DIR down"
echo "  To uninstall: curl -fsSL https://get.infrapilot.sh | bash -s -- --uninstall        (add --purge to delete data)"
echo ""