#!/usr/bin/env bash
# test-network.sh — spin up anchor + 3 peers, connect them, exchange messages + DMs.
# Usage: ./test-network.sh
# Requires: go, nc, jq, sqlite3

set -euo pipefail

# ── colours ───────────────────────────────────────────────────────────────────
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
CYAN='\033[0;36m'; MAGENTA='\033[0;35m'; BLUE='\033[0;34m'
BOLD='\033[1m'; DIM='\033[2m'; RESET='\033[0m'

ANCHOR_COLOR=$YELLOW
ALICE_COLOR=$CYAN
BOB_COLOR=$MAGENTA
CHARLIE_COLOR=$GREEN

# ── config ────────────────────────────────────────────────────────────────────
ANCHOR_PORT=19339
ALICE_IPC=19337
BOB_IPC=19340
CHARLIE_IPC=19343
NETWORK_NAME="test-$(date +%s)"
DATA_ROOT="/tmp/waste-test"
rm -rf "$DATA_ROOT"
mkdir -p "$DATA_ROOT/bin"

# Seed per-peer share directories with dummy files.
mkdir -p "$DATA_ROOT/alice/share" "$DATA_ROOT/bob/share" "$DATA_ROOT/charlie/share"
echo "alice's notes"       > "$DATA_ROOT/alice/share/notes.txt"
dd if=/dev/urandom bs=1K count=64 2>/dev/null | base64 > "$DATA_ROOT/alice/share/photo.jpg.b64"
dd if=/dev/urandom bs=1K count=128 2>/dev/null | base64 > "$DATA_ROOT/bob/share/archive.tar.b64"
echo -e "file1.mp3\nfile2.mp3" > "$DATA_ROOT/bob/share/playlist.m3u"
echo "charlie's doc"       > "$DATA_ROOT/charlie/share/document.txt"
echo "#!/bin/sh"           > "$DATA_ROOT/charlie/share/script.sh"

# Kill any leftover processes from a previous run holding our fixed ports.
for port in "$ANCHOR_PORT" "$ALICE_IPC" "$BOB_IPC" "$CHARLIE_IPC"; do
    pid=$(lsof -ti tcp:"$port" 2>/dev/null || true)
    [ -n "$pid" ] && kill -9 $pid 2>/dev/null || true
done
# Wait until all ports are actually free before proceeding.
for port in "$ANCHOR_PORT" "$ALICE_IPC" "$BOB_IPC" "$CHARLIE_IPC"; do
    n=0
    while lsof -ti tcp:"$port" >/dev/null 2>&1; do
        sleep 0.1
        n=$(( n + 1 ))
        [ "$n" -gt 30 ] && echo -e "${RED}port ${port} still in use after 3s${RESET}" >&2 && break
    done
done

# ── cleanup ───────────────────────────────────────────────────────────────────
PIDS=()
cleanup() {
    echo ""
    echo -e "${DIM}── cleaning up ──────────────────────────────────────────${RESET}"
    for pid in "${PIDS[@]:-}"; do
        kill "$pid" 2>/dev/null || true
    done
    wait 2>/dev/null || true
    echo -e "${DIM}data left at: ${DATA_ROOT}${RESET}"
    echo -e "${DIM}done.${RESET}"
}
trap cleanup EXIT INT TERM

# ── helpers ───────────────────────────────────────────────────────────────────
log() {
    local color="$1"; shift
    local label="$1"; shift
    echo -e "${color}${BOLD}[${label}]${RESET} $*"
}

# Pretty-print a raw IPC JSON line
pretty() {
    local color="$1"
    local label="$2"
    local line="$3"
    local type
    type=$(echo "$line" | jq -r '.type // "?"' 2>/dev/null)
    case "$type" in
        state_snapshot)
            local id nick
            id=$(echo "$line" | jq -r '.local_peer.id[:16]' 2>/dev/null | sed 's/.\{4\}/& /g')
            nick=$(echo "$line" | jq -r '.local_peer.alias' 2>/dev/null)
            echo -e "${color}${BOLD}[${label}]${RESET} ${DIM}state_snapshot${RESET} id=${CYAN}${id}${RESET} alias=${BOLD}${nick}${RESET}"
            ;;
        session_ready)
            local pid nick
            pid=$(echo "$line" | jq -r 'if .peer_id then .peer_id[:16] else "?" end' 2>/dev/null | sed 's/.\{4\}/& /g')
            nick=$(echo "$line" | jq -r '.nick // "?"' 2>/dev/null)
            echo -e "${color}${BOLD}[${label}]${RESET} ${GREEN}✓ session_ready${RESET} peer=${BOLD}${nick}${RESET} (${DIM}${pid}${RESET})"
            ;;
        peer_connected)
            local nick
            nick=$(echo "$line" | jq -r '.peer.alias // "?"' 2>/dev/null)
            echo -e "${color}${BOLD}[${label}]${RESET} ${GREEN}→ peer_connected${RESET} ${BOLD}${nick}${RESET}"
            ;;
        peer_disconnected)
            local pid
            pid=$(echo "$line" | jq -r 'if .peer_id then .peer_id[:8] else "?" end' 2>/dev/null)
            echo -e "${color}${BOLD}[${label}]${RESET} ${RED}← peer_disconnected${RESET} ${DIM}${pid}${RESET}"
            ;;
        message_received)
            local from body room to_field
            from=$(echo "$line" | jq -r '.message.from[:8] // "?"' 2>/dev/null)
            body=$(echo "$line" | jq -r '.message.text // ""' 2>/dev/null)
            room=$(echo "$line" | jq -r '.message.room // ""' 2>/dev/null)
            to_field=$(echo "$line" | jq -r '.message.to // ""' 2>/dev/null)
            if [ -n "$to_field" ]; then
                local to_short
                to_short=$(echo "$to_field" | cut -c1-8)
                echo -e "${color}${BOLD}[${label}]${RESET} ${BLUE}📨 DM${RESET} ${DIM}<${from}…→${to_short}…>${RESET} ${body}"
            else
                echo -e "${color}${BOLD}[${label}]${RESET} ${BOLD}💬 #${room}${RESET} ${DIM}<${from}…>${RESET} ${body}"
            fi
            ;;
        incoming_file)
            local name size xid
            name=$(echo "$line" | jq -r '.offer.name // "?"' 2>/dev/null)
            size=$(echo "$line" | jq -r '.offer.size // 0' 2>/dev/null)
            xid=$(echo "$line" | jq -r '.offer.xid[:8] // "?"' 2>/dev/null)
            echo -e "${color}${BOLD}[${label}]${RESET} ${YELLOW}↓ incoming_file${RESET} ${BOLD}${name}${RESET} (${size}B) xid=${DIM}${xid}…${RESET}"
            ;;
        file_progress)
            local xid rx total
            xid=$(echo "$line" | jq -r '.transfer_id[:8] // "?"' 2>/dev/null)
            rx=$(echo "$line" | jq -r '.bytes_received // 0' 2>/dev/null)
            total=$(echo "$line" | jq -r '.total_bytes // 0' 2>/dev/null)
            echo -e "${color}${BOLD}[${label}]${RESET} ${DIM}file_progress${RESET} xid=${xid}… ${rx}/${total}B"
            ;;
        file_complete)
            local xid path
            xid=$(echo "$line" | jq -r '.transfer_id[:8] // "?"' 2>/dev/null)
            path=$(echo "$line" | jq -r '.path // "?"' 2>/dev/null)
            echo -e "${color}${BOLD}[${label}]${RESET} ${GREEN}✓ file_complete${RESET} xid=${DIM}${xid}…${RESET} → ${BOLD}${path}${RESET}"
            ;;
        error)
            local msg
            msg=$(echo "$line" | jq -r '.error_message // .' 2>/dev/null)
            echo -e "${color}${BOLD}[${label}]${RESET} ${RED}✗ error:${RESET} ${msg}"
            ;;
        *)
            echo -e "${color}${BOLD}[${label}]${RESET} ${DIM}${line}${RESET}"
            ;;
    esac
}

# Subscribe to a peer's IPC events in the background, pretty-printing each line.
subscribe() {
    local color="$1"
    local label="$2"
    local port="$3"
    (
        set +e
        nc 127.0.0.1 "$port" 2>/dev/null | while IFS= read -r line; do
            pretty "$color" "$label" "$line"
        done
    ) &
    PIDS+=($!)
}

# Send a single JSON command to a peer's IPC port (fire-and-forget).
ipc() {
    local port="$1"
    local json="$2"
    echo "$json" | nc -q 0 127.0.0.1 "$port" >/dev/null 2>&1 || true
}

# Query get_state and return a single field via jq from the state_snapshot line.
# Retries for up to 3 seconds in case the daemon is still initialising.
# Usage: peer_field <ipc_port> <jq_expr>
peer_field() {
    local port="$1"
    local expr="$2"
    local result=""
    local n=0
    while [ -z "$result" ] || [ "$result" = "null" ]; do
        result=$(echo '{"type":"get_state"}' \
            | timeout 2 nc 127.0.0.1 "$port" 2>/dev/null \
            | grep '"type":"state_snapshot"' \
            | head -1 \
            | jq -r "$expr" 2>/dev/null || true)
        n=$(( n + 1 ))
        [ "$n" -ge 10 ] && break
        [ -z "$result" ] || [ "$result" = "null" ] && sleep 0.3
    done
    echo "$result"
}

# Wait until a TCP port accepts connections (up to N seconds).
wait_port() {
    local port="$1"
    local label="$2"
    local n=0
    while ! nc -z 127.0.0.1 "$port" 2>/dev/null; do
        sleep 0.1
        n=$(( n + 1 ))
        if [ "$n" -gt 80 ]; then
            echo -e "${RED}timeout waiting for ${label} on :${port}${RESET}" >&2
            exit 1
        fi
    done
    return 0
}

# ── build ─────────────────────────────────────────────────────────────────────
echo ""
echo -e "${BOLD}waste-go network test${RESET}"
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"
echo -e "${DIM}network : ${BOLD}${NETWORK_NAME}${RESET}"
echo -e "${DIM}data    : ${DATA_ROOT}${RESET}"
echo ""

echo -e "${DIM}building binaries…${RESET}"
go build -o "$DATA_ROOT/bin/waste-anchor" ./cmd/anchor
go build -o "$DATA_ROOT/bin/waste-daemon" ./cmd/daemon
echo -e "${GREEN}✓ built${RESET}"
echo ""

# ── anchor ────────────────────────────────────────────────────────────────────
log "$ANCHOR_COLOR" "anchor" "starting on :${ANCHOR_PORT}"
"$DATA_ROOT/bin/waste-anchor" -bind "127.0.0.1:${ANCHOR_PORT}" \
    2> >(while IFS= read -r l; do echo -e "${ANCHOR_COLOR}${DIM}[anchor] ${l}${RESET}"; done) &
PIDS+=($!)
wait_port "$ANCHOR_PORT" "anchor"
log "$ANCHOR_COLOR" "anchor" "ready"
echo ""

# ── daemons ───────────────────────────────────────────────────────────────────
ANCHOR_URL="ws://127.0.0.1:${ANCHOR_PORT}/ws"

log "$ALICE_COLOR" "alice" "starting daemon (ipc :${ALICE_IPC})"
"$DATA_ROOT/bin/waste-daemon" \
    -alias alice \
    -data-dir "$DATA_ROOT/alice" \
    -ipc-port "$ALICE_IPC" \
    -anchor "$ANCHOR_URL" \
    -share-dir "$DATA_ROOT/alice/share" \
    2> >(tee "$DATA_ROOT/alice/daemon.log" | while IFS= read -r l; do echo -e "${ALICE_COLOR}${DIM}[alice]  ${l}${RESET}"; done) &
PIDS+=($!)

log "$BOB_COLOR" "bob" "starting daemon (ipc :${BOB_IPC})"
"$DATA_ROOT/bin/waste-daemon" \
    -alias bob \
    -data-dir "$DATA_ROOT/bob" \
    -ipc-port "$BOB_IPC" \
    -anchor "$ANCHOR_URL" \
    -share-dir "$DATA_ROOT/bob/share" \
    2> >(tee "$DATA_ROOT/bob/daemon.log" | while IFS= read -r l; do echo -e "${BOB_COLOR}${DIM}[bob]    ${l}${RESET}"; done) &
PIDS+=($!)

log "$CHARLIE_COLOR" "charlie" "starting daemon (ipc :${CHARLIE_IPC})"
"$DATA_ROOT/bin/waste-daemon" \
    -alias charlie \
    -data-dir "$DATA_ROOT/charlie" \
    -ipc-port "$CHARLIE_IPC" \
    -anchor "$ANCHOR_URL" \
    -share-dir "$DATA_ROOT/charlie/share" \
    2> >(tee "$DATA_ROOT/charlie/daemon.log" | while IFS= read -r l; do echo -e "${CHARLIE_COLOR}${DIM}[charlie]${l}${RESET}"; done) &
PIDS+=($!)

wait_port "$ALICE_IPC"   "alice"
wait_port "$BOB_IPC"     "bob"
wait_port "$CHARLIE_IPC" "charlie"
echo ""

# ── join network ──────────────────────────────────────────────────────────────
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"
echo -e "joining all peers to network: ${BOLD}${NETWORK_NAME}${RESET}"
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"

JOIN=$(printf '{"type":"join_network","network_name":"%s"}' "$NETWORK_NAME")
ipc "$ALICE_IPC"   "$JOIN"
sleep 0.2
ipc "$BOB_IPC"     "$JOIN"
sleep 0.2
ipc "$CHARLIE_IPC" "$JOIN"
sleep 0.3

# ── resolve peer IDs ──────────────────────────────────────────────────────────
# Peer IDs are network-scoped (derived identity), so we query AFTER joining.
ALICE_ID=$(peer_field "$ALICE_IPC"   '.local_peer.id')
BOB_ID=$(peer_field   "$BOB_IPC"     '.local_peer.id')
CHARLIE_ID=$(peer_field "$CHARLIE_IPC" '.local_peer.id')

log "$ANCHOR_COLOR" "ids" "alice   = ${ALICE_ID:0:16}…"
log "$ANCHOR_COLOR" "ids" "bob     = ${BOB_ID:0:16}…"
log "$ANCHOR_COLOR" "ids" "charlie = ${CHARLIE_ID:0:16}…"
echo ""

# ── subscribe ─────────────────────────────────────────────────────────────────
echo -e "${DIM}subscribing to IPC event streams…${RESET}"
subscribe "$ALICE_COLOR"   "alice  " "$ALICE_IPC"
subscribe "$BOB_COLOR"     "bob    " "$BOB_IPC"
subscribe "$CHARLIE_COLOR" "charlie" "$CHARLIE_IPC"

echo -e "${DIM}waiting for ICE / DataChannel setup (up to 10s)…${RESET}"
sleep 6

# ── YAW/2.1 forward-secrecy check ────────────────────────────────────────────
# The daemon logs "2.1 FS offer" when forward-secret signaling was negotiated,
# or "2.0 fallback offer" if the peer didn't respond to the ekey in time.
# We verify by counting FS vs fallback lines in the daemon log files.
echo ""
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"
echo -e "YAW/2.1 forward-secret signaling check"
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"

fs_ok=0
fs_fail=0
for peer_data in "$DATA_ROOT/alice" "$DATA_ROOT/bob" "$DATA_ROOT/charlie"; do
    logfile="$peer_data/daemon.log"
    if [ -f "$logfile" ]; then
        ok=$(grep -c "2\.1 FS offer" "$logfile" 2>/dev/null || true)
        fail=$(grep -c "2\.0 fallback" "$logfile" 2>/dev/null || true)
        fs_ok=$(( fs_ok + ok ))
        fs_fail=$(( fs_fail + fail ))
    fi
done

# Also check stderr captured above (it was piped to terminal; count from the
# variable output buffer if possible — or just note the result from logs).
# Simpler: re-check via a short daemon log we write below.
# For now just print what we know from the test run output.
if [ "$fs_fail" -eq 0 ]; then
    echo -e "  ${GREEN}✓ all sessions negotiated YAW/2.1 (forward-secret)${RESET}"
else
    echo -e "  ${YELLOW}⚠ ${fs_fail} session(s) fell back to YAW/2.0 (non-FS)${RESET}"
fi

# ── group chat ────────────────────────────────────────────────────────────────
echo ""
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"
echo -e "group chat  (#general)"
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"

ipc "$ALICE_IPC"   '{"type":"send_message","room":"general","body":"hey everyone, alice here"}'
sleep 0.4
ipc "$BOB_IPC"     '{"type":"send_message","room":"general","body":"hi alice! bob checking in"}'
sleep 0.4
ipc "$CHARLIE_IPC" '{"type":"send_message","room":"general","body":"charlie online. the mesh works!"}'
sleep 0.4
ipc "$ALICE_IPC"   '{"type":"send_message","room":"general","body":"nice — three-way chat over WebRTC 🎉"}'
sleep 0.4
ipc "$BOB_IPC"     '{"type":"send_message","room":"general","body":"no central server, fully peer-to-peer"}'
sleep 0.4

# ── direct messages ───────────────────────────────────────────────────────────
# DMs set the "to" field to the recipient's full hex peer ID.
# The room field is used as "dm:<recipient_id>" by convention so both sides
# can query their DM history for the same conversation.
echo ""
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"
echo -e "direct messages"
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"

# Alice → Bob (privately)
ipc "$ALICE_IPC" "$(jq -cn \
    --arg to  "$BOB_ID" \
    --arg room "dm:$BOB_ID" \
    --arg body "hey bob, this is just between us" \
    '{"type":"send_message","room":$room,"body":$body,"to":$to}')"
sleep 0.4

# Bob replies to Alice
ipc "$BOB_IPC" "$(jq -cn \
    --arg to  "$ALICE_ID" \
    --arg room "dm:$ALICE_ID" \
    --arg body "got it alice, loud and clear 🤫" \
    '{"type":"send_message","room":$room,"body":$body,"to":$to}')"
sleep 0.4

# Charlie → Alice (privately)
ipc "$CHARLIE_IPC" "$(jq -cn \
    --arg to  "$ALICE_ID" \
    --arg room "dm:$ALICE_ID" \
    --arg body "alice, can you hear me? charlie dm test" \
    '{"type":"send_message","room":$room,"body":$body,"to":$to}')"
sleep 0.4

# Alice replies to Charlie
ipc "$ALICE_IPC" "$(jq -cn \
    --arg to  "$CHARLIE_ID" \
    --arg room "dm:$CHARLIE_ID" \
    --arg body "yes charlie, DMs working perfectly!" \
    '{"type":"send_message","room":$room,"body":$body,"to":$to}')"
sleep 0.4

# ── file list check ───────────────────────────────────────────────────────────
echo ""
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"
echo -e "file listing"
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"

alice_nets=$(echo '{"type":"get_state"}' | timeout 2 nc 127.0.0.1 "$ALICE_IPC" 2>/dev/null \
    | grep '"type":"state_snapshot"' | head -1 \
    | jq -r '[.networks[].network_name] | join(", ")' 2>/dev/null || echo "?")
echo -e "${DIM}  alice networks: [${alice_nets:-none}]${RESET}"

for peer_ipc in "$ALICE_IPC" "$BOB_IPC" "$CHARLIE_IPC"; do
    if [ "$peer_ipc" = "$ALICE_IPC" ]; then peer_name="alice"
    elif [ "$peer_ipc" = "$BOB_IPC" ]; then peer_name="bob"
    else peer_name="charlie"; fi
    raw=$(echo '{"type":"get_file_list"}' \
        | nc -q 2 127.0.0.1 "$peer_ipc" 2>/dev/null || true)
    result=$(echo "$raw" | grep '"type":"file_list"' | head -1 || true)
    if [ -n "$result" ]; then
        count=$(echo "$result" | jq '.files | length' 2>/dev/null || echo "?")
        files=$(echo "$result" | jq -r '[.files[].name] | join("  ")' 2>/dev/null || true)
        echo -e "  ${BOLD}${peer_name}${RESET}: ${count} file(s) — ${files}"
    else
        types=$(echo "$raw" | jq -r '.type' 2>/dev/null | tr '\n' ',' | sed 's/,$//' || true)
        echo -e "  ${RED}${peer_name}: no file_list response${RESET} ${DIM}(got: ${types:-nothing})${RESET}"
    fi
done
sleep 0.5

# ── file transfer ────────────────────────────────────────────────────────────
echo ""
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"
echo -e "file transfer (alice → bob: notes.txt)"
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"

ipc "$ALICE_IPC" "$(jq -cn --arg peer "$BOB_ID" --arg path "notes.txt" \
    '{"type":"send_file","peer_id":$peer,"path":$path}')"

# Wait up to 10s for bob to receive the file (glob over downloads-* dir)
received=0
for i in $(seq 1 50); do
    sleep 0.2
    if ls "$DATA_ROOT/bob/downloads-"*/notes.txt 2>/dev/null | grep -q .; then
        received=1; break
    fi
done

if [ "$received" = "1" ]; then
    rx_file=$(ls "$DATA_ROOT/bob/downloads-"*/notes.txt 2>/dev/null | head -1)
    orig_sha=$(sha256sum "$DATA_ROOT/alice/share/notes.txt" | awk '{print $1}')
    rx_sha=$(sha256sum "$rx_file" | awk '{print $1}')
    if [ "$orig_sha" = "$rx_sha" ]; then
        echo -e "  ${GREEN}✓ notes.txt received by bob, sha256 matches${RESET}"
    else
        echo -e "  ${RED}✗ notes.txt received but sha256 mismatch!${RESET}"
        echo -e "    orig: $orig_sha"
        echo -e "    got:  $rx_sha"
    fi
else
    echo -e "  ${RED}✗ notes.txt not received by bob within 10s${RESET}"
fi
sleep 0.5

# ── leave ─────────────────────────────────────────────────────────────────────
echo ""
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"
echo -e "charlie leaves the network"
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"
ipc "$CHARLIE_IPC" '{"type":"leave_network"}'
sleep 1

# ── persistence check ─────────────────────────────────────────────────────────
echo ""
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"
echo -e "verifying persistence (SQLite)"
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"
for peer in alice bob charlie; do
    db=$(ls "$DATA_ROOT/$peer/messages-"*.db 2>/dev/null | head -1 || true)
    if [ -n "$db" ] && [ -f "$db" ]; then
        total=$(sqlite3 "$db" "SELECT COUNT(*) FROM messages;" 2>/dev/null || echo "?")
        group=$(sqlite3 "$db" "SELECT COUNT(*) FROM messages WHERE room='general';" 2>/dev/null || echo "?")
        dms=$(sqlite3 "$db" "SELECT COUNT(*) FROM messages WHERE room LIKE 'dm:%';" 2>/dev/null || echo "?")
        peers_count=$(sqlite3 "$db" "SELECT COUNT(*) FROM peers;" 2>/dev/null || echo "?")
        echo -e "  ${BOLD}${peer}${RESET}: ${total} total (${group} group, ${dms} DM), ${peers_count} known peers"
    else
        echo -e "  ${RED}${peer}: no messages-*.db found${RESET}"
    fi
done

echo ""
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"
echo -e "${GREEN}${BOLD}test complete${RESET}"
echo -e "${DIM}────────────────────────────────────────────────────────${RESET}"
sleep 0.5