#!/bin/bash
# Agent Registration Script
# Validates and registers a new agent in Vault

set -e

VAULT_ADDR="${VAULT_ADDR:-https://127.0.0.1:8200}"
export VAULT_SKIP_VERIFY=true

usage() {
    echo "Usage: $0 -i <agent_id> -r <role> -t <tier> -o <owner> -v <version>"
    echo ""
    echo "Options:"
    echo "  -i  Agent ID (lowercase, alphanumeric with dashes)"
    echo "  -r  Role: observer|operator|builder|executor|architect"
    echo "  -t  Tier: 0-4"
    echo "  -o  Owner (human email or 'system')"
    echo "  -v  Version (semver: x.y.z)"
    echo ""
    echo "Environment:"
    echo "  VAULT_TOKEN  Required for registration"
    exit 1
}

while getopts "i:r:t:o:v:h" opt; do
    case $opt in
        i) AGENT_ID="$OPTARG" ;;
        r) ROLE="$OPTARG" ;;
        t) TIER="$OPTARG" ;;
        o) OWNER="$OPTARG" ;;
        v) VERSION="$OPTARG" ;;
        h) usage ;;
        *) usage ;;
    esac
done

# Validate required params
[[ -z "$AGENT_ID" || -z "$ROLE" || -z "$TIER" || -z "$OWNER" || -z "$VERSION" ]] && usage
[[ -z "$VAULT_TOKEN" ]] && echo "Error: VAULT_TOKEN not set" && exit 1

# Validate agent_id format
if [[ ! "$AGENT_ID" =~ ^[a-z0-9-]+$ ]]; then
    echo "Error: agent_id must be lowercase alphanumeric with dashes"
    exit 1
fi

# Validate role
VALID_ROLES="observer operator builder executor architect"
if [[ ! " $VALID_ROLES " =~ " $ROLE " ]]; then
    echo "Error: role must be one of: $VALID_ROLES"
    exit 1
fi

# Validate tier
if [[ ! "$TIER" =~ ^[0-4]$ ]]; then
    echo "Error: tier must be 0-4"
    exit 1
fi

# Validate version (semver)
if [[ ! "$VERSION" =~ ^[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
    echo "Error: version must be semver (x.y.z)"
    exit 1
fi

# Map role to tier and validate consistency
declare -A ROLE_TIER_MAP=(
    ["observer"]=0
    ["operator"]=1
    ["builder"]=2
    ["executor"]=3
    ["architect"]=4
)

EXPECTED_TIER="${ROLE_TIER_MAP[$ROLE]}"
if [[ "$TIER" -ne "$EXPECTED_TIER" ]]; then
    echo "Warning: role '$ROLE' typically maps to tier $EXPECTED_TIER, but tier $TIER was specified"
fi

# Define allowed/forbidden actions based on tier
case $TIER in
    0)
        ALLOWED='["read_docs","read_inventory","read_logs","generate_plan"]'
        FORBIDDEN='["ssh","create_vm","modify_vm","delete_vm","run_ansible","run_terraform","write_secrets","execute_shell"]'
        ;;
    1)
        ALLOWED='["read_docs","read_inventory","read_logs","generate_plan","ssh_sandbox","create_vm_sandbox","run_ansible_sandbox","run_terraform_sandbox"]'
        FORBIDDEN='["ssh_prod","ssh_staging","create_vm_prod","create_vm_staging","run_ansible_prod","run_terraform_prod","write_secrets","modify_baseline"]'
        ;;
    2)
        ALLOWED='["read_docs","read_inventory","read_logs","generate_plan","ssh_sandbox","create_vm_sandbox","run_ansible_sandbox","run_terraform_sandbox","modify_frameworks","create_templates"]'
        FORBIDDEN='["ssh_prod","create_vm_prod","run_ansible_prod","run_terraform_prod","modify_blessed_baseline","direct_prod_access"]'
        ;;
    3)
        ALLOWED='["read_docs","read_inventory","read_logs","generate_plan","ssh_sandbox","ssh_staging","ssh_prod_controlled","create_vm_sandbox","create_vm_staging","run_ansible_all","run_terraform_all"]'
        FORBIDDEN='["unbounded_root","wide_scope_apply","skip_recording","modify_governance"]'
        ;;
    4)
        ALLOWED='["read_all","propose_policy","propose_baseline","request_blessing","emergency_response"]'
        FORBIDDEN='["self_approve","self_escalate","bypass_audit"]'
        ;;
esac

# Set TTL based on tier (higher tier = shorter TTL)
TTL_MAP=(3600 1800 1800 900 900)
TTL=${TTL_MAP[$TIER]}

# Confidence threshold (higher tier = higher threshold required)
CONF_MAP=(0.7 0.75 0.8 0.85 0.9)
CONFIDENCE=${CONF_MAP[$TIER]}

TIMESTAMP=$(date -u +"%Y-%m-%dT%H:%M:%SZ")

echo "Registering agent: $AGENT_ID"
echo "  Role: $ROLE (Tier $TIER)"
echo "  Owner: $OWNER"
echo "  Version: $VERSION"
echo "  TTL: ${TTL}s"
echo "  Confidence threshold: $CONFIDENCE"

# Register in Vault
docker exec -e VAULT_TOKEN="$VAULT_TOKEN" -e VAULT_ADDR="$VAULT_ADDR" vault \
    vault kv put "secret/agents/$AGENT_ID" \
    agent_id="$AGENT_ID" \
    agent_role="$ROLE" \
    owner="$OWNER" \
    version="$VERSION" \
    tier="$TIER" \
    input_contract="secret/docs/schemas/task-request" \
    output_contract="secret/docs/schemas/agent-output" \
    allowed_side_effects="$ALLOWED" \
    forbidden_actions="$FORBIDDEN" \
    confidence_reporting=true \
    confidence_threshold="$CONFIDENCE" \
    ttl_seconds="$TTL" \
    status="registered" \
    created_at="$TIMESTAMP" \
    last_active="$TIMESTAMP" \
    compliant_runs=0 \
    consecutive_compliant=0 \
    violations=0

echo ""
echo "Agent registered successfully."
echo ""
echo "To generate credentials for this agent:"
echo "  vault read auth/approle/role/tier${TIER}-agent/role-id"
echo "  vault write -f auth/approle/role/tier${TIER}-agent/secret-id"