#!/usr/bin/env python3
"""neon-startup.py — Staged neon-64gb service startup with health checks.
   
Each stage waits for its dependencies before proceeding.
A stage failure logs an error but doesn't cascade-kill other stages.
"""

import subprocess, time, sys, os

HOST = "allaun@100.92.88.64"
SSH_OPTS = "-o ConnectTimeout=15 -o ServerAliveInterval=10"
SSH = f"ssh {SSH_OPTS} {HOST}"

def ssh(cmd: str, timeout: int = 30) -> tuple[int, str, str]:
    r = subprocess.run(f"{SSH} '{cmd}'", capture_output=True, text=True,
                       timeout=timeout, shell=True,
                       env={**os.environ, "PATH": "/usr/bin:/bin"})
    return r.returncode, r.stdout.strip(), r.stderr.strip()

def run(cmd: str, timeout: int = 30) -> int:
    rc, _, _ = ssh(cmd, timeout)
    return rc

def health(url: str) -> str:
    rc, out, _ = ssh(f"curl -s -o /dev/null -w '%{{http_code}}' http://localhost:{url}", timeout=10)
    return out if out else "down"

def wait_for(url: str, label: str, timeout: int = 60):
    deadline = time.time() + timeout
    while time.time() < deadline:
        h = health(url)
        if h in ("200", "302", "307"):
            print(f"  ✅ {label} ({h})")
            return True
        time.sleep(3)
    h = health(url)
    print(f"  ⚠️  {label} not healthy after {timeout}s (last: {h})")
    return False

def log(s): print(s)

# ── Stage 1: Infrastructure ─────────────────────────────────────────

log("=== Stage 1: Postgres ===")
run("podman start arxiv-pg 2>/dev/null", timeout=10)
time.sleep(5)
if run("pg_isready -h localhost", timeout=10) == 0:
    log("  ✅ Postgres ready")
else:
    log("  ❌ Postgres failed")

log("=== Stage 2: Redis ===")
run("podman rm -f redis 2>/dev/null; podman run -d --name redis --network host docker.io/library/redis:7-alpine", timeout=15)
wait_for("6379", "Redis", timeout=15)

# ── Stage 3: Database setup ─────────────────────────────────────────

log("=== Stage 3: Database ===")
PSQL = "PATH=$HOME/.nix-profile/bin:$PATH PGPASSWORD=postgres /home/allaun/.nix-profile/bin/psql -h localhost -U postgres"
run(f'{PSQL} -tc "SELECT 1 FROM pg_database WHERE datname='"'appflowy'"'" | grep -q 1 || '
    f'{PSQL} -c "CREATE DATABASE appflowy"', timeout=10)
run(f'{PSQL} -d appflowy -c "ALTER DATABASE appflowy SET search_path TO auth, public"', timeout=10)
run(f'{PSQL} -d appflowy -c "DROP SCHEMA IF EXISTS auth CASCADE; CREATE SCHEMA auth"', timeout=10)
log("  ✅ Database ready")

# ── Stage 4: GoTrue migrations ──────────────────────────────────────

log("=== Stage 4: GoTrue migrations ===")
MIG_IMAGE = "docker.io/appflowyinc/gotrue:0.15.1"
result = subprocess.run(
    ["podman", "run", "--rm", MIG_IMAGE, "ls", "/migrations/"],
    capture_output=True, text=True, timeout=15
)
files = sorted([f for f in result.stdout.strip().split("\n") if f.endswith(".up.sql")])

# Batch all SQL into ONE connection to avoid overwhelming postgres
all_sql = ""
for fname in files:
    r = subprocess.run(
        ["podman", "run", "--rm", MIG_IMAGE, "cat", f"/migrations/{fname}"],
        capture_output=True, text=True, timeout=15
    )
    all_sql += r.stdout.replace("{{ index .Options \"Namespace\" }}", "auth") + "\n"

# Run all migrations in one batch via SSH
proc = subprocess.Popen(
    f"{SSH} '{PSQL} -d appflowy -v ON_ERROR_STOP=0'",
    shell=True, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE
)
stdout, stderr = proc.communicate(input=all_sql.encode(), timeout=120)
log(f"  ✅ {len(files)} migrations attempted (batched)")

# Mark all versions done
all_versions = "','".join(f.split("_")[0] for f in files)
run(f'{PSQL} -d appflowy -c "INSERT INTO auth.schema_migrations (version) SELECT unnest(array[\'{all_versions}\']) ON CONFLICT DO NOTHING"', timeout=30)

# Pre-create tables that GoTrue migrations fail to create due to ordering bugs
run(f'{PSQL} -d appflowy -c "CREATE TABLE IF NOT EXISTS auth.one_time_tokens (id uuid PRIMARY KEY, user_id uuid, token_type text NOT NULL, token_hash text NOT NULL, expires_at timestamptz NOT NULL, created_at timestamptz NOT NULL DEFAULT now())"', timeout=10)

# Sync to public schema for GoTrue binary
run(f'{PSQL} -d appflowy -c "DROP TABLE IF EXISTS public.schema_migrations; CREATE TABLE public.schema_migrations AS SELECT * FROM auth.schema_migrations"', timeout=10)
log("  Migration tracking synced")

# ── Stage 5: GoTrue service ─────────────────────────────────────────

log("=== Stage 5: GoTrue ===")
run("podman rm -f appflowy_gotrue_1 2>/dev/null", timeout=5)
run(
    'podman run -d --name appflowy_gotrue_1 --network host '
    '-e GOTRUE_ADMIN_EMAIL=admin@researchstack.info '
    '-e GOTRUE_ADMIN_PASSWORD=admin123 '
    '-e GOTRUE_DISABLE_SIGNUP=false '
    '-e GOTRUE_SITE_URL=https://researchstack.info/appflowy '
    '-e GOTRUE_URI_ALLOW_LIST=** '
    '-e GOTRUE_JWT_SECRET=jwt_ '
    '-e GOTRUE_JWT_EXP=3600 '
    '-e GOTRUE_DB_DRIVER=postgres '
    '-e API_EXTERNAL_URL=https://researchstack.info/appflowy/gotrue '
    '-e DATABASE_URL="postgres://postgres:postgres@localhost:5432/appflowy" '
    '-e PORT=9999 '
    '-e GOTRUE_MAILER_AUTOCONFIRM=true '
    '-e GOTRUE_LOG_LEVEL=info '
    '-e REDIS_ENABLED=true '
    '-e REDIS_URL=redis://localhost:6379 '
    'docker.io/appflowyinc/gotrue:0.15.1', timeout=20)
wait_for("9999/health", "GoTrue", timeout=30)

# ── Stage 6: AppFloyo Cloud ─────────────────────────────────────────

log("=== Stage 6: AppFloyo Cloud ===")
run("podman rm -f appflowy_cloud 2>/dev/null", timeout=5)
run(
    'podman run -d --name appflowy_cloud --network host '
    '-e APPFLOWY_DATABASE_URL="postgres://postgres:postgres@localhost:5432/appflowy?sslmode=disable&options=-c%20search_path=public" '
    '-e APPFLOWY_ENVIRONMENT=production '
    '-e APPFLOWY_REDIS_URI=redis://localhost:6379 '
    '-e APPFLOWY_GOTRUE_JWT_SECRET=jwt_ '
    '-e APPFLOWY_GOTRUE_JWT_EXP=3600 '
    '-e APPFLOWY_GOTRUE_BASE_URL=http://localhost:9999 '
    '-e APPFLOWY_BASE_URL=https://researchstack.info/appflowy '
    '-e APPFLOWY_WEB_URL=https://researchstack.info/appflowy '
    '-e APPFLOWY_ENCRYPTION_KEY=sk_fix_me_32_chars_abcdef1234567890 '
    '-e APPFLOWY_JWT_SECRET=jwt_ '
    '-e APPFLOWY_CREATE_BUCKET=false '
    '-e APPFLOWY_S3_ENABLED=false '
    '-e APPFLOWY_SEARCH_ENABLED=false '
    '-e RUST_LOG=info '
    'docker.io/appflowyinc/appflowy_cloud:latest', timeout=20)
wait_for("8000/api/health", "AppFloyo", timeout=45)

# ── Stage 7: Authentik SSO ──────────────────────────────────────────

log("=== Stage 7: Authentik ===")
run("podman start authentik-server 2>/dev/null", timeout=10)
run("podman start authentik-worker 2>/dev/null", timeout=10)
wait_for("30001/", "Authentik", timeout=45)

# ── Stage 8: Supporting services ────────────────────────────────────

log("=== Stage 8: Other services ===")
for svc, port, label in [("couchdb", "5984/", "CouchDB"),
                          ("homarr-web", "7575/auth/login", "Homarr"),
                          ("www-static", "8082/", "www-static")]:
    run(f"podman start {svc} 2>/dev/null", timeout=10)
    wait_for(port, label, timeout=30)

# ── Final report ────────────────────────────────────────────────────

log("\n=== Service Summary ===")
for name, url in [("Postgres (pg)", "5432 (pg_isready)"),
                  ("Redis", "6379 (redis-cli ping)"),
                  ("GoTrue", "9999/health"),
                  ("AppFloyo", "8000/api/health"),
                  ("Authentik", "30001/"),
                  ("CouchDB", "5984/"),
                  ("Homarr", "7575/auth/login")]:
    if "pg" in name.lower():
        rc = run("pg_isready -h localhost", timeout=5)
        log(f"  {'✅' if rc==0 else '❌'} {name}")
    else:
        h = health(url.split("/")[0] + "/" + "/".join(url.split("/")[1:]))
        log(f"  {'✅' if h in ('200','302','307') else '❌'} {name} ({h})")