#!/usr/bin/env bash
#
# HeroCtl Installer
# Usage: curl -sfL https://get.heroctl.com/install.sh | HEROCTL_TOKEN=<token> sh -
#
# Options (environment variables):
#   HEROCTL_TOKEN=<base64>               # join token (contains server addresses, auto-configures agent)
#   HEROCTL_VERSION=v0.1.0               # specific version (default: latest)
#   HEROCTL_MODE=agent                   # "server", "agent", or "both" (default: agent when token provided)
#   HEROCTL_SERVER_ADDR=10.0.0.1:8080    # server address (manual, used when no token)
#
# Examples:
#   # Agent with join token (1 command!)
#   curl -sfL https://get.heroctl.com/install.sh | HEROCTL_TOKEN=eyJzZXJ2... bash -
#
#   # Single node (server + agent)
#   curl -sfL https://get.heroctl.com/install.sh | HEROCTL_MODE=both bash -
#
#   # Server only
#   curl -sfL https://get.heroctl.com/install.sh | HEROCTL_MODE=server sh -
#
set -eu

HEROCTL_TOKEN="${HEROCTL_TOKEN:-}"
HEROCTL_VERSION="${HEROCTL_VERSION:-latest}"
HEROCTL_SERVER_ADDR="${HEROCTL_SERVER_ADDR:-localhost:8080}"

# If token is provided, default to agent mode
if [ -n "$HEROCTL_TOKEN" ]; then
    HEROCTL_MODE="${HEROCTL_MODE:-agent}"
else
    HEROCTL_MODE="${HEROCTL_MODE:-both}"
fi

INSTALL_DIR="/usr/local/bin"
CONFIG_DIR="/etc/heroctl"
DATA_DIR="/var/lib/heroctl"
SYSTEMD_DIR="/etc/systemd/system"

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

log()  { echo -e "${GREEN}==>${NC} $1"; }
warn() { echo -e "${YELLOW}==>${NC} $1"; }
err()  { echo -e "${RED}ERROR:${NC} $1"; exit 1; }

# --- Checks ---

if [ "$(id -u)" -ne 0 ]; then
    err "This script must be run as root (use sudo)"
fi

ARCH=$(uname -m)
case "$ARCH" in
    x86_64)  ARCH="amd64" ;;
    aarch64) ARCH="arm64" ;;
    armv7l)  ARCH="arm" ;;
    *)       err "Unsupported architecture: $ARCH" ;;
esac

OS=$(uname -s | tr '[:upper:]' '[:lower:]')
if [ "$OS" != "linux" ]; then
    err "This installer is for Linux only. OS detected: $OS"
fi

# --- Decode join token ---

TOKEN_SERVERS=""
TOKEN_TLS="false"

if [ -n "$HEROCTL_TOKEN" ]; then
    log "Decoding join token..."
    TOKEN_JSON=$(echo "$HEROCTL_TOKEN" | base64 -d 2>/dev/null) || err "Invalid token (base64 decode failed)"

    # Extract servers array using python3 or simple parsing
    if command -v python3 &>/dev/null; then
        TOKEN_SERVERS=$(echo "$TOKEN_JSON" | python3 -c "import sys,json; d=json.load(sys.stdin); print(' '.join(d.get('servers',[])))" 2>/dev/null)
        TOKEN_TLS=$(echo "$TOKEN_JSON" | python3 -c "import sys,json; print(str(json.load(sys.stdin).get('tls',False)).lower())" 2>/dev/null)
    else
        # Fallback: basic parsing with grep/sed
        TOKEN_SERVERS=$(echo "$TOKEN_JSON" | grep -o '"servers":\[[^]]*\]' | grep -o '[0-9.]*:[0-9]*' | tr '\n' ' ')
    fi

    if [ -z "$TOKEN_SERVERS" ]; then
        err "Token contains no server addresses"
    fi

    HEROCTL_SERVER_ADDR=$(echo "$TOKEN_SERVERS" | awk '{print $1}')
    log "  Servers: $TOKEN_SERVERS"
    log "  TLS: $TOKEN_TLS"
fi

# --- Check container runtime ---

DRIVER="docker"
if command -v podman &>/dev/null; then
    DRIVER="podman"
    log "Detected Podman"
elif command -v docker &>/dev/null; then
    DRIVER="docker"
    log "Detected Docker"
else
    warn "No container runtime found. Install Docker or Podman first."
    warn "  Docker: curl -fsSL https://get.docker.com | sh"
    warn "  Podman: apt install -y podman"
fi

# --- Install binaries ---

log "Installing HeroCtl ${HEROCTL_VERSION} (${OS}/${ARCH}) — mode: ${HEROCTL_MODE}"

# Try to download from server first, then check local files
DOWNLOAD_SUCCESS=false
if [ -n "$TOKEN_SERVERS" ]; then
    FIRST_SERVER=$(echo "$TOKEN_SERVERS" | awk '{print $1}')
    PROTO="http"
    [ "$TOKEN_TLS" = "true" ] && PROTO="https"
    DOWNLOAD_URL="${PROTO}://${FIRST_SERVER}"

    log "  Downloading binaries from server ${FIRST_SERVER}..."
    for bin in heroctl-agent heroctl; do
        if curl -sfL -k "${DOWNLOAD_URL}/v1/agent/binary/${bin}?arch=${ARCH}" -o "${INSTALL_DIR}/${bin}" 2>/dev/null; then
            chmod +x "${INSTALL_DIR}/${bin}"
            log "  Downloaded ${bin}"
            DOWNLOAD_SUCCESS=true
        fi
    done
fi

# Fallback: local files
if [ "$DOWNLOAD_SUCCESS" = "false" ]; then
    for bin in heroctl-server heroctl-agent heroctl; do
        if [ -f "./${bin}" ]; then
            cp "./${bin}" "${INSTALL_DIR}/${bin}"
            chmod +x "${INSTALL_DIR}/${bin}"
            log "  Installed ${bin} (local)"
        elif [ -f "./bin/${bin}" ]; then
            cp "./bin/${bin}" "${INSTALL_DIR}/${bin}"
            chmod +x "${INSTALL_DIR}/${bin}"
            log "  Installed ${bin} (local)"
        elif [ -f "./bin/linux/${bin}" ]; then
            cp "./bin/linux/${bin}" "${INSTALL_DIR}/${bin}"
            chmod +x "${INSTALL_DIR}/${bin}"
            log "  Installed ${bin} (local)"
        else
            if [ "$HEROCTL_MODE" = "agent" ] && [ "$bin" = "heroctl-server" ]; then
                continue  # agent mode doesn't need server binary
            fi
            warn "  Binary ${bin} not found. Build first: make build-linux"
        fi
    done
fi

# --- Create directories ---

mkdir -p "${CONFIG_DIR}" "${CONFIG_DIR}/ingress.d"
mkdir -p "${DATA_DIR}/server" "${DATA_DIR}/agent"
log "Created directories"

# --- Write configs ---

if [ "$HEROCTL_MODE" = "server" ] || [ "$HEROCTL_MODE" = "both" ]; then
    if [ ! -f "${CONFIG_DIR}/server.yaml" ]; then
        cat > "${CONFIG_DIR}/server.yaml" <<YAML
# HeroCtl Server Configuration
bind_addr: "0.0.0.0"
http_port: 8080
grpc_port: 9090
data_dir: "${DATA_DIR}/server"
log_level: "info"
dns_port: 8600
tls_enabled: false
ingress_enabled: false
ingress_binary_path: "${INSTALL_DIR}/traefik"
ingress_config_dir: "${CONFIG_DIR}/ingress.d"
YAML
        log "Created ${CONFIG_DIR}/server.yaml"
    fi
fi

if [ "$HEROCTL_MODE" = "agent" ] || [ "$HEROCTL_MODE" = "both" ]; then
    # Build agent config from token or defaults
    if [ -n "$TOKEN_SERVERS" ]; then
        # Token-based config: multi-server with all addresses
        SERVER_ADDRS_YAML=""
        for srv in $TOKEN_SERVERS; do
            SERVER_ADDRS_YAML="${SERVER_ADDRS_YAML}\n  - \"${srv}\""
        done

        cat > "${CONFIG_DIR}/agent.yaml" <<YAML
# HeroCtl Agent Configuration (auto-generated from join token)
server_addr: "${HEROCTL_SERVER_ADDR}"
server_addrs:$(echo -e "$SERVER_ADDRS_YAML")
driver: "${DRIVER}"
data_dir: "${DATA_DIR}/agent"
log_level: "info"
YAML
        log "Created ${CONFIG_DIR}/agent.yaml (from join token)"
    elif [ ! -f "${CONFIG_DIR}/agent.yaml" ]; then
        cat > "${CONFIG_DIR}/agent.yaml" <<YAML
# HeroCtl Agent Configuration
server_addr: "${HEROCTL_SERVER_ADDR}"
driver: "${DRIVER}"
data_dir: "${DATA_DIR}/agent"
log_level: "info"
YAML
        log "Created ${CONFIG_DIR}/agent.yaml"
    fi
fi

# --- Install systemd service files ---

if [ "$HEROCTL_MODE" = "server" ] || [ "$HEROCTL_MODE" = "both" ]; then
    cat > "${SYSTEMD_DIR}/heroctl-server.service" <<'UNIT'
[Unit]
Description=HeroCtl Server - Workload Orchestrator
Documentation=https://github.com/heroio/heroctl
Wants=network-online.target
After=network-online.target

[Service]
Type=simple
EnvironmentFile=-/etc/heroctl/server.env
ExecStart=/usr/local/bin/heroctl-server --config /etc/heroctl/server.yaml
KillSignal=SIGINT
KillMode=process
Restart=on-failure
RestartSec=5
LimitNOFILE=65536
OOMScoreAdjust=-1000

[Install]
WantedBy=multi-user.target
UNIT
    log "Installed heroctl-server.service"
fi

if [ "$HEROCTL_MODE" = "agent" ] || [ "$HEROCTL_MODE" = "both" ]; then
    cat > "${SYSTEMD_DIR}/heroctl-agent.service" <<'UNIT'
[Unit]
Description=HeroCtl Agent - Workload Orchestrator Worker
Documentation=https://github.com/heroio/heroctl
Wants=network-online.target
After=network-online.target docker.service podman.service

[Service]
Type=simple
EnvironmentFile=-/etc/heroctl/agent.env
ExecStart=/usr/local/bin/heroctl-agent --config /etc/heroctl/agent.yaml
KillSignal=SIGINT
KillMode=process
Restart=on-failure
RestartSec=5
LimitNOFILE=65536
OOMScoreAdjust=-1000

[Install]
WantedBy=multi-user.target
UNIT
    log "Installed heroctl-agent.service"
fi

systemctl daemon-reload

# --- Install Traefik (server mode only) ---

if [ "$HEROCTL_MODE" = "server" ] || [ "$HEROCTL_MODE" = "both" ]; then
    if ! command -v traefik &>/dev/null && [ ! -f "${INSTALL_DIR}/traefik" ]; then
        log "Downloading Traefik v3.6..."
        TRAEFIK_URL="https://github.com/traefik/traefik/releases/download/v3.6.13/traefik_v3.6.13_linux_${ARCH}.tar.gz"
        curl -fsSL "$TRAEFIK_URL" | tar xz -C "${INSTALL_DIR}" traefik
        chmod +x "${INSTALL_DIR}/traefik"
        log "Installed Traefik to ${INSTALL_DIR}/traefik"
    else
        log "Traefik already installed"
    fi
fi

# --- Enable and start services ---

if [ "$HEROCTL_MODE" = "server" ] || [ "$HEROCTL_MODE" = "both" ]; then
    systemctl enable heroctl-server
    log "Enabled heroctl-server (auto-start on boot)"
fi

if [ "$HEROCTL_MODE" = "agent" ] || [ "$HEROCTL_MODE" = "both" ]; then
    systemctl enable heroctl-agent
    log "Enabled heroctl-agent (auto-start on boot)"
fi

# --- Auto-start if token provided (agent ready to go) ---

if [ -n "$HEROCTL_TOKEN" ]; then
    log "Starting agent (join token provided)..."
    systemctl start heroctl-agent
    sleep 3
    if systemctl is-active --quiet heroctl-agent; then
        log "Agent started successfully!"
    else
        warn "Agent failed to start. Check: journalctl -u heroctl-agent -f"
    fi
fi

# --- Done ---

echo ""
echo "============================================"
echo "  HeroCtl installed successfully!"
echo "============================================"
echo ""
echo "  Config:  ${CONFIG_DIR}/"
echo "  Data:    ${DATA_DIR}/"
echo "  Logs:    journalctl -u heroctl-agent -f"
echo ""

if [ -n "$HEROCTL_TOKEN" ]; then
    echo "  Agent is running and connected to:"
    for srv in $TOKEN_SERVERS; do
        echo "    - ${srv}"
    done
    echo ""
    echo "  Verify: heroctl node list"
else
    if [ "$HEROCTL_MODE" = "server" ] || [ "$HEROCTL_MODE" = "both" ]; then
        echo "  Edit ${CONFIG_DIR}/server.yaml, then:"
        echo "    sudo systemctl start heroctl-server"
        echo ""
    fi
    if [ "$HEROCTL_MODE" = "agent" ] || [ "$HEROCTL_MODE" = "both" ]; then
        echo "  Edit ${CONFIG_DIR}/agent.yaml, then:"
        echo "    sudo systemctl start heroctl-agent"
        echo ""
    fi
fi
echo ""