#!/bin/bash
set -euo pipefail

# ================================================================
# Restore script for prod Matrix instances (yc-playground → yandex-prod)
# ================================================================
# Reads a backup created by backup-yc-playground.sh and restores
# it to the new yandex-prod cluster with ESS chart.
#
# Usage: ./restore-prod.sh <name> <backup-dir>
# Example: ./restore-prod.sh t0rt1k backups/matrix-t0rt1k-20260614-161207/
# ================================================================

readonly YC_KUBECONFIG="${KUBECONFIG:-/home/mrt0rtikize/infra/yandex-prod/kubeconfig}"
readonly K="${KUBECTL:-kubectl} --kubeconfig ${YC_KUBECONFIG}"

readonly CNPG_NS="cnpg"
readonly CNPG_POD="shared-pg-1"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'
log()  { echo -e "${GREEN}[$(date +%H:%M:%S)]${NC} $*"; }
warn() { echo -e "${YELLOW}[$(date +%H:%M:%S)] WARN:${NC} $*"; }
err()  { echo -e "${RED}[$(date +%H:%M:%S)] ERROR:${NC} $*"; }

# -------------------------------------------------------------------
# Parse args
# -------------------------------------------------------------------
if [ $# -ne 2 ]; then
    echo "Usage: $0 <name> <backup-dir>"
    echo ""
    echo "  name:       t0rt1k | roglog | uretra"
    echo "  backup-dir: path to backup directory from backup-yc-playground.sh"
    echo ""
    echo "Example: $0 t0rt1k backups/matrix-t0rt1k-20260614-161207/"
    exit 1
fi

readonly NAME="$1"
readonly NS="matrix-${NAME}"
readonly DB_SYNAPSE="synapse_${NAME}"
readonly DB_MAS="mas_${NAME}"
readonly BACKUP_DIR="$2"

# -------------------------------------------------------------------
# Prerequisites
# -------------------------------------------------------------------
log "=== Restoring ${NS} from ${BACKUP_DIR} ==="

if [ ! -d "${BACKUP_DIR}" ]; then
    err "Backup directory not found: ${BACKUP_DIR}"
    exit 1
fi

for f in dump-synapse.sql dump-mas.sql secret-chat-matrix.yaml secret-mas.yaml secret-chat-postgresql.yaml; do
    if [ ! -f "${BACKUP_DIR}/${f}" ]; then
        err "Missing: ${f}"
        exit 1
    fi
done

if ! ${K} get ns "${NS}" >/dev/null 2>&1; then
    err "Namespace ${NS} not found. Deploy the ESS app and sync first."
    exit 1
fi

if ! ${K} -n "${CNPG_NS}" get pod "${CNPG_POD}" >/dev/null 2>&1; then
    err "CNPG pod ${CNPG_POD} not found."
    exit 1
fi

log "Prerequisites OK."

# -------------------------------------------------------------------
# Step 1: Stop Synapse + MAS
# -------------------------------------------------------------------
log "=== Step 1: Stopping Synapse + MAS ==="

${K} -n "${NS}" scale deploy -l "app.kubernetes.io/component=matrix-server" --replicas=0 2>/dev/null || \
    ${K} -n "${NS}" scale sts "${NS}-synapse-main" --replicas=0 2>/dev/null || \
    ${K} -n "${NS}" scale deploy chat-matrix --replicas=0 2>/dev/null || \
    warn "Could not scale Synapse."

${K} -n "${NS}" scale deploy -l "app.kubernetes.io/component=matrix-authentication" --replicas=0 2>/dev/null || \
    ${K} -n "${NS}" scale deploy "${NS}-matrix-authentication-service" --replicas=0 2>/dev/null || \
    ${K} -n "${NS}" scale deploy mas --replicas=0 2>/dev/null || \
    warn "Could not scale MAS."

log "Waiting for pods to terminate..."
${K} -n "${NS}" wait --for=delete pod -l "app.kubernetes.io/component=matrix-server" --timeout=120s 2>/dev/null || warn "Synapse may still be terminating."
${K} -n "${NS}" wait --for=delete pod -l "app.kubernetes.io/component=matrix-authentication" --timeout=120s 2>/dev/null || warn "MAS may still be terminating."
${K} -n "${NS}" wait --for=delete pod -l "app=mas" --timeout=60s 2>/dev/null || true

log "Synapse + MAS stopped."

# -------------------------------------------------------------------
# Step 2: Clean schemas
# -------------------------------------------------------------------
log "=== Step 2: Cleaning database schemas ==="

log "Reading PG credentials..."
SYNAPSE_PW=$(${K} get secret pg-creds -n "${NS}" -o jsonpath='{.data.synapse}' 2>/dev/null | base64 -d)
MAS_PW=$(${K} get secret pg-creds -n "${NS}" -o jsonpath='{.data.mas}' 2>/dev/null | base64 -d)

log "Wiping ${DB_SYNAPSE} schema..."
${K} exec -n "${CNPG_NS}" "${CNPG_POD}" -- env PGPASSWORD="${SYNAPSE_PW}" \
    psql -U "${DB_SYNAPSE}" -d "${DB_SYNAPSE}" -h localhost -c \
    "DROP SCHEMA public CASCADE; CREATE SCHEMA public; GRANT ALL ON SCHEMA public TO ${DB_SYNAPSE};" 2>/dev/null || true

log "Wiping ${DB_MAS} schema..."
${K} exec -n "${CNPG_NS}" "${CNPG_POD}" -- env PGPASSWORD="${MAS_PW}" \
    psql -U "${DB_MAS}" -d "${DB_MAS}" -h localhost -c \
    "DROP SCHEMA public CASCADE; CREATE SCHEMA public; GRANT ALL ON SCHEMA public TO ${DB_MAS};" 2>/dev/null || true

log "Schemas cleaned."

# -------------------------------------------------------------------
# Step 3: Restore PostgreSQL dumps
# -------------------------------------------------------------------
log "=== Step 3: Restoring PostgreSQL dumps ==="

log "Restoring Synapse database..."
# Old chart used 'matrix' role — create it as superuser to suppress OWNER TO errors
${K} exec -n "${CNPG_NS}" "${CNPG_POD}" -- psql -U postgres -c \
    "DO \$\$ BEGIN IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = 'matrix') THEN CREATE ROLE matrix; END IF; END \$\$;" 2>/dev/null || true
${K} exec -n "${CNPG_NS}" "${CNPG_POD}" -- psql -U postgres -c \
    "GRANT matrix TO ${DB_SYNAPSE};" 2>/dev/null || true

${K} exec -i -n "${CNPG_NS}" "${CNPG_POD}" -- env PGPASSWORD="${SYNAPSE_PW}" \
    psql -U "${DB_SYNAPSE}" -d "${DB_SYNAPSE}" -h localhost < "${BACKUP_DIR}/dump-synapse.sql"
log "Synapse database restored."

log "Restoring MAS database..."
# Old chart used 'mas' role — create it as superuser to suppress OWNER TO errors
${K} exec -n "${CNPG_NS}" "${CNPG_POD}" -- psql -U postgres -c \
    "DO \$\$ BEGIN IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = 'mas') THEN CREATE ROLE mas; END IF; END \$\$;" 2>/dev/null || true
${K} exec -n "${CNPG_NS}" "${CNPG_POD}" -- psql -U postgres -c \
    "GRANT mas TO ${DB_MAS};" 2>/dev/null || true
${K} exec -i -n "${CNPG_NS}" "${CNPG_POD}" -- env PGPASSWORD="${MAS_PW}" \
    psql -U "${DB_MAS}" -d "${DB_MAS}" -h localhost < "${BACKUP_DIR}/dump-mas.sql"
log "MAS database restored."

# -------------------------------------------------------------------
# Step 4: Construct and apply the ESS -generated secret
# -------------------------------------------------------------------
log "=== Step 4: Constructing ESS generated secret ==="

# Read existing auto-generated secret (has keys the old chart didn't: ECDSA, LiveKit, Hookshot, etc.)
EXISTING_KEY_ECDSA=""
EXISTING_KEY_LIVEKIT=""
EXISTING_KEY_HOOKSHOT_REG=""
EXISTING_KEY_HOOKSHOT_PASS=""
EXISTING_KEY_SYNAPSE_EXTRA=""
EXISTING_KEY_PG_ADMIN=""

if ${K} -n "${NS}" get secret "${NS}-generated" >/dev/null 2>&1; then
    log "Reading existing auto-generated secret for non-migrated keys..."
    EXISTING_KEY_ECDSA=$(${K} -n "${NS}" get secret "${NS}-generated" -o jsonpath='{.data.MAS_ECDSA_PRIME256V1_PRIVATE_KEY}' 2>/dev/null || echo "")
    EXISTING_KEY_LIVEKIT=$(${K} -n "${NS}" get secret "${NS}-generated" -o jsonpath='{.data.ELEMENT_CALL_LIVEKIT_SECRET}' 2>/dev/null || echo "")
    EXISTING_KEY_HOOKSHOT_REG=$(${K} -n "${NS}" get secret "${NS}-generated" -o jsonpath='{.data.HOOKSHOT_REGISTRATION}' 2>/dev/null || echo "")
    EXISTING_KEY_HOOKSHOT_PASS=$(${K} -n "${NS}" get secret "${NS}-generated" -o jsonpath='{.data.HOOKSHOT_RSA_PASSKEY}' 2>/dev/null || echo "")
    EXISTING_KEY_SYNAPSE_EXTRA=$(${K} -n "${NS}" get secret "${NS}-generated" -o jsonpath='{.data.SYNAPSE_EXTRA}' 2>/dev/null || echo "")
    EXISTING_KEY_PG_ADMIN=$(${K} -n "${NS}" get secret "${NS}-generated" -o jsonpath='{.data.POSTGRES_ADMIN_PASSWORD}' 2>/dev/null || echo "")
    log "Preserved non-migrated keys from existing secret."
else
    warn "No existing ${NS}-generated secret found. Some keys may be missing."
fi

# Read old secrets from backup
log "Reading old secrets from backup..."
OLD_SIGNING_KEY=$(python3 -c "
import yaml, sys
with open('${BACKUP_DIR}/secret-chat-matrix.yaml') as f:
    s = yaml.safe_load(f.read())
print(s['data']['signing.key'])
" 2>/dev/null)

OLD_MACAROON=$(python3 -c "
import yaml, sys
with open('${BACKUP_DIR}/secret-chat-matrix.yaml') as f:
    s = yaml.safe_load(f.read())
print(s['data']['macaroon_secret_key'])
" 2>/dev/null)

OLD_REGISTRATION=$(python3 -c "
import yaml, sys
with open('${BACKUP_DIR}/secret-chat-matrix.yaml') as f:
    s = yaml.safe_load(f.read())
print(s['data']['registration_shared_secret'])
" 2>/dev/null)

OLD_MAS_ENCRYPTION=$(python3 -c "
import yaml, sys
with open('${BACKUP_DIR}/secret-mas.yaml') as f:
    s = yaml.safe_load(f.read())
print(s['data']['encryption-key'])
" 2>/dev/null)

OLD_MAS_RSA=$(python3 -c "
import yaml, sys
with open('${BACKUP_DIR}/secret-mas.yaml') as f:
    s = yaml.safe_load(f.read())
print(s['data']['signing-key.pem'])
" 2>/dev/null)

OLD_MAS_SHARED=$(python3 -c "
import yaml, sys
with open('${BACKUP_DIR}/secret-mas.yaml') as f:
    s = yaml.safe_load(f.read())
print(s['data']['shared-secret'])
" 2>/dev/null)

OLD_SYNAPSE_PG_PW=$(python3 -c "
import yaml, sys
with open('${BACKUP_DIR}/secret-chat-postgresql.yaml') as f:
    s = yaml.safe_load(f.read())
print(s['data']['password'])
" 2>/dev/null)

OLD_MAS_DB_PW=$(python3 -c "
import yaml, sys
with open('${BACKUP_DIR}/secret-mas.yaml') as f:
    s = yaml.safe_load(f.read())
print(s['data']['mas-db-password'])
" 2>/dev/null)

# Build the new secret
log "Applying generated secret with migrated keys..."

PYTHON_BODY=$(cat <<'PYEOF'
import yaml, sys, base64

# Read existing secret structure from stdin (if any)
existing = {}
try:
    existing = yaml.safe_load(sys.stdin)
except:
    pass

data = existing.get('data', {}) if existing else {}

# Overwrite with old keys (these are base64 encoded already)
data['SYNAPSE_SIGNING_KEY'] = sys.argv[1] if sys.argv[1] else data.get('SYNAPSE_SIGNING_KEY','')
data['SYNAPSE_MACAROON'] = sys.argv[2] if sys.argv[2] else data.get('SYNAPSE_MACAROON','')
data['SYNAPSE_REGISTRATION_SHARED_SECRET'] = sys.argv[3] if sys.argv[3] else data.get('SYNAPSE_REGISTRATION_SHARED_SECRET','')
data['MAS_ENCRYPTION_SECRET'] = sys.argv[4] if sys.argv[4] else data.get('MAS_ENCRYPTION_SECRET','')
data['MAS_RSA_PRIVATE_KEY'] = sys.argv[5] if sys.argv[5] else data.get('MAS_RSA_PRIVATE_KEY','')

data['MAS_SYNAPSE_SHARED_SECRET'] = sys.argv[6] if sys.argv[6] else data.get('MAS_SYNAPSE_SHARED_SECRET','')
data['POSTGRES_SYNAPSE_PASSWORD'] = sys.argv[7] if sys.argv[7] else data.get('POSTGRES_SYNAPSE_PASSWORD','')
data['POSTGRES_MATRIX_AUTHENTICATION_SERVICE_PASSWORD'] = sys.argv[8] if sys.argv[8] else data.get('POSTGRES_MATRIX_AUTHENTICATION_SERVICE_PASSWORD','')

# Preserve existing non-migrated keys if provided
if sys.argv[9]: data['MAS_ECDSA_PRIME256V1_PRIVATE_KEY'] = sys.argv[9]
if sys.argv[10]: data['ELEMENT_CALL_LIVEKIT_SECRET'] = sys.argv[10]
if sys.argv[11]: data['HOOKSHOT_REGISTRATION'] = sys.argv[11]
if sys.argv[12]: data['HOOKSHOT_RSA_PASSKEY'] = sys.argv[12]
if sys.argv[13]: data['SYNAPSE_EXTRA'] = sys.argv[13]
if sys.argv[14]: data['POSTGRES_ADMIN_PASSWORD'] = sys.argv[14]

secret = {
    'apiVersion': 'v1',
    'kind': 'Secret',
    'metadata': {
        'name': f'{sys.argv[15]}-generated',
        'namespace': sys.argv[16],
    },
    'type': 'Opaque',
    'data': data,
}

yaml.dump(secret, sys.stdout, default_flow_style=False)
PYEOF
)

# Get existing secret for merge
${K} -n "${NS}" get secret "${NS}-generated" -o yaml 2>/dev/null | \
    python3 -c "${PYTHON_BODY}" \
    "${OLD_SIGNING_KEY}" "${OLD_MACAROON}" "${OLD_REGISTRATION}" \
    "${OLD_MAS_ENCRYPTION}" "${OLD_MAS_RSA}" "${OLD_MAS_SHARED}" \
    "${OLD_SYNAPSE_PG_PW}" "${OLD_MAS_DB_PW}" \
    "${EXISTING_KEY_ECDSA}" "${EXISTING_KEY_LIVEKIT}" "${EXISTING_KEY_HOOKSHOT_REG}" \
    "${EXISTING_KEY_HOOKSHOT_PASS}" "${EXISTING_KEY_SYNAPSE_EXTRA}" "${EXISTING_KEY_PG_ADMIN}" \
    "${NS}" "${NS}" | \
    ${K} replace --force -f - 2>/dev/null || \
    ${K} create -f - 2>/dev/null

log "Generated secret applied with migrated keys from old chart."

# -------------------------------------------------------------------
# Step 5: Scale Synapse to 1
# -------------------------------------------------------------------
log "=== Step 5: Starting Synapse ==="

${K} -n "${NS}" scale sts -l "app.kubernetes.io/component=matrix-server" --replicas=1 2>/dev/null || \
    ${K} -n "${NS}" scale sts "${NS}-synapse-main" --replicas=1 2>/dev/null || \
    ${K} -n "${NS}" scale deploy chat-matrix --replicas=1 2>/dev/null

log "Waiting for Synapse..."
${K} -n "${NS}" wait --for=condition=ready pod -l "app.kubernetes.io/component=matrix-server" --timeout=300s 2>/dev/null || warn "Synapse not ready yet."

# -------------------------------------------------------------------
# Step 6: Restore media files
# -------------------------------------------------------------------
log "=== Step 6: Restoring media files ==="

if [ -f "${BACKUP_DIR}/synapse-media.tar.gz" ]; then
    SYNAPSE_POD=$( ${K} -n "${NS}" get pods -l "app.kubernetes.io/component=matrix-server" -o jsonpath='{.items[?(@.status.phase=="Running")].metadata.name}' 2>/dev/null)
    if [ -z "${SYNAPSE_POD}" ]; then
        warn "No running Synapse pod for media restore."
    else
        log "Copying media to Synapse pod: ${SYNAPSE_POD}"
        ${K} cp "${BACKUP_DIR}/synapse-media.tar.gz" "${NS}/${SYNAPSE_POD}:/tmp/synapse-media.tar.gz"
        ${K} exec -n "${NS}" "${SYNAPSE_POD}" -- tar xzf /tmp/synapse-media.tar.gz -C /media/
        ${K} exec -n "${NS}" "${SYNAPSE_POD}" -- rm /tmp/synapse-media.tar.gz
        log "Media files restored."
    fi
else
    warn "synapse-media.tar.gz not found in backup."
fi

# -------------------------------------------------------------------
# Step 7: Scale MAS to 1
# -------------------------------------------------------------------
log "=== Step 7: Starting MAS ==="

${K} -n "${NS}" scale deploy -l "app.kubernetes.io/component=matrix-authentication" --replicas=1 2>/dev/null || \
    ${K} -n "${NS}" scale deploy "${NS}-matrix-authentication-service" --replicas=1 2>/dev/null || \
    ${K} -n "${NS}" scale deploy mas --replicas=1 2>/dev/null

log "Waiting for MAS..."
${K} -n "${NS}" wait --for=condition=ready pod -l "app.kubernetes.io/component=matrix-authentication" --timeout=120s 2>/dev/null || warn "MAS not ready yet."

# -------------------------------------------------------------------
# Summary
# -------------------------------------------------------------------
cat <<SUMMARY

${GREEN}========================================${NC}
${GREEN}  Restore Complete: ${NS}${NC}
${GREEN}========================================${NC}

  Verification:
  1. Check Synapse logs:     ${K} logs -n ${NS} -l "app.kubernetes.io/component=matrix-server" --tail=20
  2. Test login:             https://chat.${NAME}.tech
  3. Federation tester:      https://federationtester.matrix.org/?server_name=${NAME}.tech

SUMMARY