#!/bin/bash
# Sei Labs: fast node entry for Sei mainnet (pacific-1).
#
#   curl -fsSL https://dev-labs.seistream.app/install.sh | bash
#   curl -fsSL https://dev-labs.seistream.app/install.sh | bash -s -- --home /data/sei --jobs 16
#
# What it does:
#   1. downloads every part of the latest artifact in parallel, checks sha256 while streaming,
#      unpacks each part straight into place (no tarball on disk);
#   2. recomputes the SHA-256 hashes that were stripped from memIAVL nodes/leaves
#      and checks every tree root against the snapshot's own commit info;
#   3. sets the two app.toml keys this artifact needs (ss-enable=false, snapshot-interval=0)
#      and persistent peers.
# It never touches priv_validator_key.json / node_key.json and refuses a non-empty data/.
set -euo pipefail

BASE=${SEI_LABS_URL:-https://dev-labs.seistream.app}
NET=mainnet
HOME_DIR=${SEI_HOME:-$HOME/.sei}
JOBS=8
SEID=${SEID:-seid}
FORCE_VERSION=0

while [ $# -gt 0 ]; do
  case "$1" in
    --home) HOME_DIR=$2; shift 2 ;;
    --jobs) JOBS=$2; shift 2 ;;
    --url) BASE=$2; shift 2 ;;
    --seid) SEID=$2; shift 2 ;;
    --force-version) FORCE_VERSION=1; shift ;;
    -h|--help) sed -n '2,16p' "$0" 2>/dev/null; exit 0 ;;
    *) echo "unknown option: $1" >&2; exit 2 ;;
  esac
done

say(){ printf "\033[1m[sei-labs]\033[0m %s\n" "$*"; }
fail(){ printf "\033[31m[sei-labs] %s\033[0m\n" "$*" >&2; exit 1; }

for t in curl zstd tar python3 sha256sum; do command -v $t >/dev/null || fail "missing tool: $t"; done
[ "$(uname -sm)" = "Linux x86_64" ] || fail "linux x86_64 only"
T0=$(date +%s)

TMP=$(mktemp -d); trap 'rm -rf "$TMP"' EXIT
curl -fsSL "$BASE/manifests/$NET/latest.json" -o "$TMP/m.json" || fail "cannot fetch manifest"
mf(){ python3 -c "import json,sys;m=json.load(open('$TMP/m.json'));print(eval(sys.argv[1]))" "$1"; }
HEIGHT=$(mf 'm["height"]'); WANT=$(mf 'm["seid"]["version"]'); SIZE=$(mf 'm["size_bytes"]'); RAW=$(mf 'm["raw_bytes"]')
say "artifact: height $HEIGHT, $(( SIZE/1073741824 )) GiB to download, seid $WANT"

# версия seid обязана совпадать: другой бинарник уже однажды испортил состояние при открытии
if command -v "$SEID" >/dev/null; then
  HAVE=$("$SEID" version 2>/dev/null | head -1 | tr -d '[:space:]')
  [ "v${HAVE#v}" = "v${WANT#v}" ] || [ $FORCE_VERSION = 1 ] || fail "seid $HAVE installed, artifact needs $WANT (use --force-version to override)"
else
  say "warning: seid not found in PATH; you must run exactly $WANT"
fi

[ -d "$HOME_DIR/data/state_commit" ] && [ -n "$(ls -A "$HOME_DIR/data/state_commit" 2>/dev/null)" ] \
  && fail "$HOME_DIR/data/state_commit is not empty; move data/ aside first"
mkdir -p "$HOME_DIR/data"
AVAIL=$(df -B1 --output=avail "$HOME_DIR" | tail -1)
NEED=$(python3 -c "import json;m=json.load(open('$TMP/m.json'));print(m['raw_bytes']+2*sum(f['size'] for f in m['stripped_files']))")
[ "$AVAIL" -gt "$NEED" ] || fail "need $((NEED/1073741824)) GiB free in $HOME_DIR, have $((AVAIL/1073741824))"

if [ ! -f "$HOME_DIR/config/config.toml" ]; then
  command -v "$SEID" >/dev/null || fail "no config in $HOME_DIR and no seid to create one"
  "$SEID" init sei-labs-node --chain-id pacific-1 --home "$HOME_DIR" >/dev/null 2>&1
  curl -fsSL "$BASE$(mf 'm["client"]["genesis"]["url"]')" -o "$HOME_DIR/config/genesis.json"
  say "created config in $HOME_DIR/config"
fi
[ "$(sha256sum < "$HOME_DIR/config/genesis.json" | cut -d' ' -f1)" = "$(mf 'm["client"]["genesis"]["sha256"]')" ] \
  || say "warning: your genesis.json differs from the one this artifact was built with"

HT=$TMP/hashtool
curl -fsSL "$BASE$(mf 'm["client"]["hashtool"]["url"]')" -o "$HT"
[ "$(sha256sum < "$HT" | cut -d' ' -f1)" = "$(mf 'm["client"]["hashtool"]["sha256"]')" ] || fail "hashtool checksum mismatch"
chmod +x "$HT"

say "1/3 downloading $(mf 'len(m["parts"])') parts with $JOBS parallel streams"
python3 - "$TMP/m.json" "$BASE" "$HOME_DIR" "$JOBS" <<'PY'
import concurrent.futures as cf, hashlib, json, os, subprocess, sys, threading, time
m = json.load(open(sys.argv[1])); base, home, jobs = sys.argv[2], sys.argv[3], int(sys.argv[4])
need = {}
for f in m["range_files"]:
    need[f["path"]] = max(need.get(f["path"], 0), int(f["size"]))
for part in m["parts"]:
    if part.get("kind") != "range":
        continue
    end = int(part["offset"]) + int(part["length"])
    need[part["path"]] = max(need.get(part["path"], 0), end)
for path, size in need.items():
    fp = os.path.join(home, path)
    os.makedirs(os.path.dirname(fp), exist_ok=True)
    with open(fp, "ab") as h:
        h.truncate(size)
done = [0]; lock = threading.Lock(); total = m["size_bytes"]; t0 = time.time()

def fetch(part):
    for attempt in range(5):
        try:
            return _fetch(part)
        except Exception as e:
            err = e; time.sleep(2 + attempt * 3)
    raise RuntimeError(f"part {part['sha256'][:12]}: {err}")

def _fetch(part):
    curl = subprocess.Popen(["curl", "-fsSL", "--retry", "3", base + part["url"]], stdout=subprocess.PIPE)
    if part["kind"] == "tar":
        sink = subprocess.Popen(["tar", "-x", "-C", home, "-f", "-"], stdin=subprocess.PIPE)
        zstd = subprocess.Popen(["zstd", "-dc"], stdin=subprocess.PIPE, stdout=sink.stdin)
        out = None
    else:
        zstd = subprocess.Popen(["zstd", "-dc"], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
        out = open(os.path.join(home, part["path"]), "r+b"); out.seek(part["offset"])
        def writer():
            for b in iter(lambda: zstd.stdout.read(8 << 20), b""): out.write(b)
        wt = threading.Thread(target=writer); wt.start()
    h = hashlib.sha256(); n = 0
    for b in iter(lambda: curl.stdout.read(4 << 20), b""):
        h.update(b); zstd.stdin.write(b); n += len(b)
        with lock: done[0] += len(b)
    zstd.stdin.close()
    if curl.wait(): raise RuntimeError("download failed")
    if zstd.wait(): raise RuntimeError("zstd failed")
    if out:
        wt.join(); size = out.tell() - part["offset"]; out.close()
        if size != part["length"]: raise RuntimeError(f"length {size} != {part['length']}")
    else:
        sink.stdin.close()
        if sink.wait(): raise RuntimeError("tar failed")
    if h.hexdigest() != part["sha256"]:
        with lock: done[0] -= n
        raise RuntimeError("sha256 mismatch")

def progress(stop):
    while not stop.wait(10):
        with lock: d = done[0]
        el = time.time() - t0
        print(f"      {d/2**30:6.1f} / {total/2**30:.1f} GiB  {d*8/el/1e6:6.0f} Mbit/s", flush=True)

stop = threading.Event(); threading.Thread(target=progress, args=(stop,), daemon=True).start()
parts = sorted(m["parts"], key=lambda p: -p["bytes"])
with cf.ThreadPoolExecutor(jobs) as ex:
    for f in cf.as_completed([ex.submit(fetch, p) for p in parts]): f.result()
stop.set()
print(f"      all {len(parts)} parts verified in {time.time()-t0:.0f} s", flush=True)
PY

say "2/3 recomputing memIAVL hashes"
MEM=$HOME_DIR/data/state_commit/memiavl
SNAP=snapshot-$(printf "%020d" "$(mf 'm["memiavl_snapshot"]')")
"$HT" roots "$MEM/$SNAP" > "$TMP/roots.txt"
refill_one(){
  local d=$1 m; m=$(basename "$d")
  [ -f "$d/nodes" ] || : > "$d/nodes"
  mv "$d/nodes" "$d/nodes.bare"; mv "$d/leaves" "$d/leaves.bare"
  local r; r=$("$HT" refill "$d/nodes.bare" "$d/leaves.bare" "$d/kvs" "$d/nodes" "$d/leaves" | awk '/^ROOT/{print $2}')
  rm -f "$d/nodes.bare" "$d/leaves.bare"
  [ -n "$r" ] || r=e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855   # пустое дерево
  local want; want=$(awk -v m="$m" '$1==m{print $2}' "$TMP/roots.txt")
  if [ "$r" = "$want" ]; then echo "      $m ok"; else echo "      $m ROOT MISMATCH: $r != $want"; return 1; fi
}
export -f refill_one; export HT TMP
# сначала самые большие деревья, по два одновременно (каждое уже грузит все ядра)
du -s "$MEM/$SNAP"/*/ | sort -rn | cut -f2 | xargs -P 2 -I{} bash -c 'refill_one "$1"' _ {} \
  || fail "tree root check failed; do not start the node on this data"
ln -sfn "$SNAP" "$MEM/current"

say "3/3 config"
[ -f "$HOME_DIR/data/priv_validator_state.json" ] || echo '{"height":"0","round":0,"step":0}' > "$HOME_DIR/data/priv_validator_state.json"
python3 - "$HOME_DIR" "$(mf 'm["client"]["persistent_peers"]')" <<'PY'
import re, sys
home, peers = sys.argv[1], sys.argv[2]
p = home + "/config/app.toml"; s = open(p).read()
s = re.sub(r'(?m)^ss-enable\s*=.*', 'ss-enable = false', s)
# только в [state-sync]: в [state-commit.flatkv] есть одноимённый ключ (giga), его не трогаем
s = re.sub(r'(?ms)(^\[state-sync\][^\[]*?^snapshot-interval\s*=\s*)\S+', r'\g<1>0', s)
open(p, "w").write(s)
p = home + "/config/config.toml"; s = open(p).read()
if peers:
    s = re.sub(r'(?m)^persistent-peers\s*=.*', f'persistent-peers = "{peers}"', s, count=1)
    s = re.sub(r'(?m)^bootstrap-peers\s*=\s*""', f'bootstrap-peers = "{peers}"', s, count=1)
s = re.sub(r'(?ms)(^\[statesync\][^\[]*?^enable\s*=\s*)true', r'\1false', s)
open(p, "w").write(s)
PY
say "done in $(( $(date +%s)-T0 )) s. Start the node:  $SEID start --home $HOME_DIR"
say "it replays the changelog to height $HEIGHT in seconds, then catches up from peers"