#!/bin/bash
# 0G Labs: fast node entry for 0G mainnet (Aristotle, chain 16661).
#
#   curl -fsSL https://labs.0g.exploreme.pro/install.sh | bash -s -- --home ~/aristotle-v1.0.6/0g-home
#   curl -fsSL https://labs.0g.exploreme.pro/install.sh | bash -s -- --home /data/0g-home --release ~/aristotle-v1.0.6 --jobs 16
#
# --home is the 0g-home directory of the official release (it has 0gchaind-home/ and geth-home/).
# What it does:
#   1. checks that your 0gchaind and geth are exactly the versions the snapshot was taken with;
#   2. downloads every part in parallel, checks sha256 while streaming and writes each part
#      straight into place (no tarball on disk);
#   3. creates node keys if the release has none, sets persistent_peers if yours are empty.
# It never overwrites priv_validator_key.json / node_key.json (creates them only if missing)
# and refuses non-empty data.
set -euo pipefail

BASE=${OG_LABS_URL:-https://labs.0g.exploreme.pro}
NET=mainnet
HOME_DIR=""
RELEASE=""
JOBS=8
FORCE_VERSION=0

while [ $# -gt 0 ]; do
  case "$1" in
    --home) HOME_DIR=$2; shift 2 ;;
    --release) RELEASE=$2; shift 2 ;;
    --jobs) JOBS=$2; shift 2 ;;
    --url) BASE=$2; shift 2 ;;
    --net) NET=$2; shift 2 ;;     # для проверок: manifests/<net>/latest.json
    --force-version) FORCE_VERSION=1; shift ;;
    -h|--help) sed -n '2,15p' "$0" 2>/dev/null; exit 0 ;;
    *) echo "unknown option: $1" >&2; exit 2 ;;
  esac
done

say(){ printf "\033[1m[0g-labs]\033[0m %s\n" "$*"; }
fail(){ printf "\033[31m[0g-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"
[ -n "$HOME_DIR" ] || { [ -d ./0g-home ] && HOME_DIR=./0g-home; } || fail "pass --home <0g-home of the aristotle release>"
HOME_DIR=$(cd "$HOME_DIR" && pwd) || fail "no such directory: $HOME_DIR"
[ -n "$RELEASE" ] || RELEASE=$(dirname "$HOME_DIR")
[ -f "$HOME_DIR/0gchaind-home/config/config.toml" ] \
  || fail "$HOME_DIR/0gchaind-home/config is missing: unpack the official release and pass its 0g-home"
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"]'); SIZE=$(mf 'm["size_bytes"]')
WCL=$(mf 'm["binaries"]["0gchaind"]'); WEL=$(mf 'm["binaries"]["geth"]'); RELNAME=$(mf 'm["binaries"]["release"]')
say "snapshot: height $HEIGHT, $(( SIZE/1073741824 )) GiB to download, $RELNAME (0gchaind $WCL, geth $WEL)"

CLB=$RELEASE/bin/0gchaind; ELB=$RELEASE/bin/geth
if [ -x "$CLB" ] && [ -x "$ELB" ]; then
  HCL=$("$CLB" version 2>/dev/null | head -1 | tr -d '[:space:]')
  HEL=$("$ELB" version 2>/dev/null | awk -F': *' '/^Version:/{print $2}')
  { [ "$HCL" = "$WCL" ] && [ "$HEL" = "$WEL" ]; } || [ $FORCE_VERSION = 1 ] \
    || fail "you have 0gchaind $HCL / geth $HEL, the snapshot needs $WCL / $WEL (use --force-version to override)"
else
  say "warning: no binaries in $RELEASE/bin; you must run exactly 0gchaind $WCL and geth $WEL"
fi

D=$HOME_DIR/0gchaind-home/data
[ -n "$(ls -A "$HOME_DIR/geth-home/geth/chaindata" 2>/dev/null)" ] && fail "$HOME_DIR/geth-home/geth/chaindata is not empty; move it aside first"
for d in application.db blockstore.db state.db; do [ -e "$D/$d" ] && fail "$D/$d exists; move data/ aside first"; done
mkdir -p "$D" "$HOME_DIR/geth-home"
AVAIL=$(df -B1 --output=avail "$HOME_DIR" | tail -1)
NEED=$(( $(mf 'm["raw_bytes"]') * 11 / 10 ))
[ "$AVAIL" -gt "$NEED" ] || fail "need $((NEED/1073741824)) GiB free in $HOME_DIR, have $((AVAIL/1073741824))"
[ "$(sha256sum < "$HOME_DIR/0gchaind-home/config/genesis.json" | cut -d' ' -f1)" = "$(mf 'm["client"]["config"]["genesis.json"]["sha256"]')" ] \
  || say "warning: your genesis.json differs from the one this snapshot was taken with"

say "1/2 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])
for f in m["range_files"]:
    fp = os.path.join(home, f["path"])
    os.makedirs(os.path.dirname(fp), exist_ok=True)
    with open(fp, "ab") as h:
        h.truncate(int(f["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/2 config"
# 0gchaind не стартует без priv_validator_key.json (BLS-подписант), а в релизе ключей нет:
# недостающие создаём через init во временном каталоге, существующие не трогаем никогда
CFG=$HOME_DIR/0gchaind-home/config
if [ ! -f "$CFG/priv_validator_key.json" ] || [ ! -f "$CFG/node_key.json" ]; then
  [ -x "$CLB" ] || fail "no keys in $CFG and no 0gchaind to create them"
  "$CLB" init labs-node --home "$TMP/keys" --chaincfg.chain-spec mainnet >/dev/null 2>&1 || fail "0gchaind init failed"
  for k in priv_validator_key.json node_key.json; do
    [ -f "$CFG/$k" ] || { install -m600 "$TMP/keys/config/$k" "$CFG/$k"; say "created $k (new keys for this node)"; }
  done
fi
python3 - "$HOME_DIR/0gchaind-home/config/config.toml" "$(mf 'm["client"]["persistent_peers"]')" <<'PY'
import re, sys
p, peers = sys.argv[1], sys.argv[2]
s = open(p).read()
if peers and re.search(r'(?m)^persistent_peers\s*=\s*""', s):
    s = re.sub(r'(?m)^persistent_peers\s*=\s*""', f'persistent_peers = "{peers}"', s, count=1)
open(p, "w").write(s)
PY
say "done in $(( $(date +%s)-T0 )) s, node at block $HEIGHT. Start geth first, then 0gchaind (from $RELEASE):"
cat <<EOF

  ./bin/geth --config geth-config.toml --datadir $HOME_DIR/geth-home --networkid 16661 \\
      --authrpc.addr 127.0.0.1 --authrpc.port 8551 --authrpc.jwtsecret jwt.hex

  ./bin/0gchaind start --home $HOME_DIR/0gchaind-home \\
      --chaincfg.kzg.trusted-setup-path=kzg-trusted-setup.json \\
      --chaincfg.engine.jwt-secret-path=jwt.hex --chaincfg.block-store-service.enabled \\
      --pruning=custom --pruning-keep-recent=100 --pruning-interval=19

EOF
say "stop geth with SIGINT and give it time: a killed geth loses its recent state and 0gchaind stalls"