#!/bin/bash
#
# Tier 0 Agent Bootstrap Script
# =============================
# Authenticates with Vault, obtains scoped token, initializes agent environment.
#
# Usage:
#   ./bootstrap.sh
#   source ./bootstrap.sh  # To export VAULT_TOKEN to current shell
#

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CONFIG_DIR="${SCRIPT_DIR}/config"
CREDS_DIR="${SCRIPT_DIR}/credentials"
LOGS_DIR="${SCRIPT_DIR}/logs"
WORKSPACE_DIR="${SCRIPT_DIR}/workspace"

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

log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
log_ok() { echo -e "${GREEN}[OK]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

# -----------------------------------------------------------------------------
# Configuration
# -----------------------------------------------------------------------------

VAULT_ADDR="${VAULT_ADDR:-https://127.0.0.1:8200}"
AGENT_CONFIG="${CONFIG_DIR}/agent.json"
APPROLE_CREDS="${CREDS_DIR}/approle.json"

# Read agent config
if [[ ! -f "${AGENT_CONFIG}" ]]; then
    log_error "Agent config not found: ${AGENT_CONFIG}"
    exit 1
fi

AGENT_ID=$(python3 -c "import json; print(json.load(open('${AGENT_CONFIG}'))['agent_id'])")
AGENT_TIER=$(python3 -c "import json; print(json.load(open('${AGENT_CONFIG}'))['tier'])")

echo ""
echo "=========================================="
echo "TIER 0 AGENT BOOTSTRAP"
echo "=========================================="
echo "Agent ID: ${AGENT_ID}"
echo "Tier: ${AGENT_TIER}"
echo "Vault: ${VAULT_ADDR}"
echo ""

# -----------------------------------------------------------------------------
# Step 1: Vault Authentication
# -----------------------------------------------------------------------------

log_info "Authenticating with Vault..."

if [[ ! -f "${APPROLE_CREDS}" ]]; then
    log_error "AppRole credentials not found: ${APPROLE_CREDS}"
    exit 1
fi

ROLE_ID=$(python3 -c "import json; print(json.load(open('${APPROLE_CREDS}'))['role_id'])")
SECRET_ID=$(python3 -c "import json; print(json.load(open('${APPROLE_CREDS}'))['secret_id'])")

# Login to Vault
LOGIN_RESPONSE=$(curl -sk -X POST \
    -d "{\"role_id\":\"${ROLE_ID}\",\"secret_id\":\"${SECRET_ID}\"}" \
    "${VAULT_ADDR}/v1/auth/approle/login")

# Check for errors
if echo "${LOGIN_RESPONSE}" | grep -q '"errors"'; then
    log_error "Vault login failed:"
    echo "${LOGIN_RESPONSE}" | python3 -m json.tool
    exit 1
fi

# Extract token
VAULT_TOKEN=$(echo "${LOGIN_RESPONSE}" | python3 -c "import sys,json; print(json.load(sys.stdin)['auth']['client_token'])")
TOKEN_TTL=$(echo "${LOGIN_RESPONSE}" | python3 -c "import sys,json; print(json.load(sys.stdin)['auth']['lease_duration'])")
TOKEN_POLICIES=$(echo "${LOGIN_RESPONSE}" | python3 -c "import sys,json; print(','.join(json.load(sys.stdin)['auth']['policies']))")

log_ok "Vault authentication successful"
echo "  Token TTL: ${TOKEN_TTL}s"
echo "  Policies: ${TOKEN_POLICIES}"

# Export token
export VAULT_TOKEN
export VAULT_ADDR

# Save token to file for agent processes
echo "${VAULT_TOKEN}" > "${CREDS_DIR}/.token"
chmod 600 "${CREDS_DIR}/.token"

# -----------------------------------------------------------------------------
# Step 2: Verify Token Permissions
# -----------------------------------------------------------------------------

log_info "Verifying token permissions..."

# Test token lookup (should work with agent-self-read)
TOKEN_INFO=$(curl -sk -H "X-Vault-Token: ${VAULT_TOKEN}" \
    "${VAULT_ADDR}/v1/auth/token/lookup-self" 2>/dev/null || echo '{"errors":["failed"]}')

if echo "${TOKEN_INFO}" | grep -q '"errors"'; then
    log_warn "Could not verify token (agent-self-read may not be attached)"
else
    log_ok "Token self-lookup successful"
fi

# Test that we CANNOT access secrets (Tier 0 should be denied)
SECRETS_TEST=$(curl -sk -H "X-Vault-Token: ${VAULT_TOKEN}" \
    "${VAULT_ADDR}/v1/secret/data/services/dragonfly" 2>/dev/null || echo '{}')

if echo "${SECRETS_TEST}" | grep -q '"data"'; then
    log_error "SECURITY VIOLATION: Tier 0 token can access secrets!"
    exit 1
else
    log_ok "Confirmed: Cannot access secrets (as expected for Tier 0)"
fi

# Test that we CANNOT access SSH (Tier 0 should be denied)
SSH_TEST=$(curl -sk -H "X-Vault-Token: ${VAULT_TOKEN}" \
    "${VAULT_ADDR}/v1/ssh/creds/sandbox-user" -X POST -d '{"ip":"10.77.10.1"}' 2>/dev/null || echo '{}')

if echo "${SSH_TEST}" | grep -q '"signed_key"'; then
    log_error "SECURITY VIOLATION: Tier 0 token can get SSH credentials!"
    exit 1
else
    log_ok "Confirmed: Cannot access SSH credentials (as expected for Tier 0)"
fi

# -----------------------------------------------------------------------------
# Step 3: Initialize Agent Environment
# -----------------------------------------------------------------------------

log_info "Initializing agent environment..."

# Export agent environment variables
export AGENT_ID
export AGENT_TIER
export AGENT_CONFIG
export AGENT_WORKSPACE="${WORKSPACE_DIR}"
export AGENT_LOGS="${LOGS_DIR}"

# Create session ID
export SESSION_ID="sess-$(date +%Y%m%d-%H%M%S)-$(openssl rand -hex 4)"
echo "${SESSION_ID}" > "${WORKSPACE_DIR}/.session_id"

# Initialize log file
LOG_FILE="${LOGS_DIR}/agent-${SESSION_ID}.log"
echo "=== Agent Session Started ===" > "${LOG_FILE}"
echo "Session ID: ${SESSION_ID}" >> "${LOG_FILE}"
echo "Agent ID: ${AGENT_ID}" >> "${LOG_FILE}"
echo "Tier: ${AGENT_TIER}" >> "${LOG_FILE}"
echo "Timestamp: $(date -u +%Y-%m-%dT%H:%M:%SZ)" >> "${LOG_FILE}"
echo "" >> "${LOG_FILE}"

log_ok "Agent environment initialized"
echo "  Session: ${SESSION_ID}"
echo "  Workspace: ${WORKSPACE_DIR}"
echo "  Logs: ${LOG_FILE}"

# -----------------------------------------------------------------------------
# Step 4: Register with Governance System
# -----------------------------------------------------------------------------

log_info "Registering with governance system..."

# Register in DragonflyDB
DRAGONFLY_CREDS=$(curl -sk -H "X-Vault-Token: $(cat /opt/vault/init-keys.json | python3 -c "import sys,json; print(json.load(sys.stdin)['root_token'])")" \
    "${VAULT_ADDR}/v1/secret/data/services/dragonfly" 2>/dev/null || echo '{}')

if echo "${DRAGONFLY_CREDS}" | grep -q '"password"'; then
    REDIS_PASS=$(echo "${DRAGONFLY_CREDS}" | python3 -c "import sys,json; print(json.load(sys.stdin)['data']['data']['password'])")

    # Register agent state
    redis-cli -p 6379 -a "${REDIS_PASS}" SET "agent:${AGENT_ID}:state" "{\"agent_id\":\"${AGENT_ID}\",\"tier\":${AGENT_TIER},\"status\":\"RUNNING\",\"session_id\":\"${SESSION_ID}\",\"started_at\":\"$(date -u +%Y-%m-%dT%H:%M:%SZ)\"}" 2>/dev/null

    # Set heartbeat
    redis-cli -p 6379 -a "${REDIS_PASS}" SET "agent:${AGENT_ID}:heartbeat" "$(date +%s)" EX 60 2>/dev/null

    log_ok "Registered with governance system"
else
    log_warn "Could not register with DragonflyDB (non-fatal)"
fi

# Register in SQLite ledger
sqlite3 /opt/agent-governance/ledger/governance.db "
INSERT OR REPLACE INTO agent_metrics (agent_id, current_tier, compliant_runs, consecutive_compliant, total_runs, updated_at)
VALUES ('${AGENT_ID}', ${AGENT_TIER}, 0, 0, 0, datetime('now'));
" 2>/dev/null || log_warn "Could not update ledger (non-fatal)"

# -----------------------------------------------------------------------------
# Done
# -----------------------------------------------------------------------------

echo ""
echo "=========================================="
echo "BOOTSTRAP COMPLETE"
echo "=========================================="
echo "Agent ${AGENT_ID} is ready (Tier ${AGENT_TIER})"
echo ""
echo "Environment variables set:"
echo "  VAULT_TOKEN  (use for Vault API calls)"
echo "  VAULT_ADDR   = ${VAULT_ADDR}"
echo "  AGENT_ID     = ${AGENT_ID}"
echo "  AGENT_TIER   = ${AGENT_TIER}"
echo "  SESSION_ID   = ${SESSION_ID}"
echo ""
echo "Next: Run the agent with ./run-agent.sh"
echo "=========================================="