//! Distributed lockstep BFT consensus over real QUIC.
//!
//! Each process is one node. Nodes derive a shared signed membership
//! deterministically from their indices, form a full QUIC mesh (joule-moq
//! transport), and run a single lockstep round — broadcast a signed proposal,
//! barrier, broadcast the per-node vote set, barrier, then commit deterministically
//! in `op_id` order via [`joule_code_collab::RoundNode`]. Every honest node prints
//! the same agreed-log hash, with no leader.
//!
//! TLS is a shared pinned self-signed cert (generate once, distribute); the BFT
//! identities are the per-node Ed25519 op-signing keys, independent of TLS.
//!
//! Usage:
//!   lockstep-node gen-cert <cert_out.der> <key_out.der>
//!   lockstep-node run <index> <addr0,addr1,...> <cert.der> <key.der>

use std::collections::{BTreeMap, BTreeSet};
use std::sync::Arc;
use std::time::Duration;

use ed25519_dalek::SigningKey;
use joule_code_cargo::{Anchor, CodingOp, Edit, SignedOp};
use joule_code_collab::{Membership, RoundNode};
use joule_code_schema::{ContentId, Tier, ALG_ED25519};
use joule_moq::transport_quic::{generate_self_signed_der, QuicEndpoint, QuicSession};
use serde::{Deserialize, Serialize};

/// One wire message in a lockstep round.
#[derive(Serialize, Deserialize)]
enum Msg {
    /// A node's signed proposal for the round.
    Propose { from: usize, op: SignedOp, prior: ContentId },
    /// A node's accept-vote set (the `op_id`s it validated).
    Vote { from: usize, accept: BTreeSet<String> },
}

/// Deterministic per-node Ed25519 identity, so every node computes the same
/// membership without exchanging keys.
fn node_key(index: usize) -> SigningKey {
    SigningKey::from_bytes(&[(100 + index) as u8; 32])
}

/// The op this node proposes — authored by its own key, so every honest node
/// validates it identically.
fn build_proposal(index: usize) -> (SignedOp, ContentId) {
    let key = node_key(index);
    let prior = ContentId("b3:GENESIS".into());
    let op = CodingOp {
        alg: ALG_ED25519,
        prior: prior.clone(),
        edits: vec![Edit {
            path: format!("node-{index}.txt"),
            anchor: Anchor::ByteRange { start: 0, end: 0 },
            replacement: format!("proposed by node {index}").into_bytes(),
        }],
        author: key.verifying_key().to_bytes(),
        intent: format!("node-{index} round-0 proposal"),
        tier: Tier::L4,
    };
    (op.sign(&key).expect("sign proposal"), prior)
}

async fn send_msg(session: &QuicSession, msg: &Msg) -> Result<(), Box<dyn std::error::Error>> {
    let body = serde_json::to_vec(msg)?;
    let mut stream = session.open_uni().await?;
    stream.write_all(&(body.len() as u32).to_be_bytes()).await?;
    stream.write_all(&body).await?;
    stream.finish()?;
    Ok(())
}

async fn recv_msg(session: &QuicSession) -> Result<Msg, Box<dyn std::error::Error>> {
    let mut stream = session.accept_uni().await?;
    let mut len = [0u8; 4];
    stream.read_exact(&mut len).await?;
    let n = u32::from_be_bytes(len) as usize;
    let mut buf = vec![0u8; n];
    stream.read_exact(&mut buf).await?;
    Ok(serde_json::from_slice(&buf)?)
}

fn log_hash(node: &RoundNode) -> String {
    let mut hasher = blake3::Hasher::new();
    for committed in node.log() {
        hasher.update(&committed.sequence.to_be_bytes());
        hasher.update(committed.op.op.op_id().expect("op_id").0.as_bytes());
    }
    hasher.finalize().to_hex().to_string()
}

#[tokio::main(flavor = "multi_thread", worker_threads = 2)]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = std::env::args().collect();
    match args.get(1).map(String::as_str) {
        Some("gen-cert") => {
            let cert_out = args.get(2).ok_or("need cert out")?;
            let key_out = args.get(3).ok_or("need key out")?;
            let (cert_der, key_der) = generate_self_signed_der(vec!["joule-node".into()])?;
            std::fs::write(cert_out, &cert_der)?;
            std::fs::write(key_out, &key_der)?;
            println!("wrote shared cert -> {cert_out}, key -> {key_out}");
            Ok(())
        }
        Some("run") => run(&args).await,
        Some("star") => run_star(&args).await,
        _ => {
            eprintln!("usage: lockstep-node gen-cert <cert_out> <key_out>");
            eprintln!("       lockstep-node run <index> <addr0,addr1,...> <cert.der> <key.der>");
            eprintln!(
                "       lockstep-node star <index> <hub_index> <addr0,addr1,...> <cert.der> <key.der>"
            );
            std::process::exit(2);
        }
    }
}

async fn run(args: &[String]) -> Result<(), Box<dyn std::error::Error>> {
    let index: usize = args.get(2).ok_or("need index")?.parse()?;
    let addrs: Vec<std::net::SocketAddr> = args
        .get(3)
        .ok_or("need addr list")?
        .split(',')
        .map(|s| s.trim().parse())
        .collect::<Result<_, _>>()?;
    let cert_path = args.get(4).ok_or("need cert path")?;
    let key_path = args.get(5).ok_or("need key path")?;
    let n = addrs.len();
    if index >= n {
        return Err("index out of range".into());
    }

    let cert_der = std::fs::read(cert_path)?;
    let key_der = std::fs::read(key_path)?;

    let membership =
        Membership::new((0..n).map(|i| node_key(i).verifying_key().to_bytes()).collect());
    let mut node = RoundNode::new(index, membership, n);
    println!("node {index}/{n}: f={} threshold={}", node.f(), node.supermajority());

    // ── Form the mesh: node i accepts from lower indices, dials higher ones. ──
    // One dual-role endpoint on 0.0.0.0:<port> — it both accepts inbound peers
    // and dials outbound ones on a single socket. Running a separate listening
    // endpoint and dialing endpoint in the same process wedges cross-machine
    // dials on these multi-homed hosts; one endpoint dials and accepts cleanly.
    let listen = std::net::SocketAddr::from(([0, 0, 0, 0], addrs[index].port()));
    let endpoint = Arc::new(QuicEndpoint::node_from_der(
        listen,
        cert_der.clone(),
        key_der,
        cert_der,
    )?);

    let accept_count = index; // nodes 0..index dial us
    let accept_handle = tokio::spawn({
        let endpoint = endpoint.clone();
        async move {
            let mut accepted = Vec::new();
            for _ in 0..accept_count {
                match endpoint.accept().await {
                    Ok(session) => {
                        eprintln!("  accepted a peer ({}/{})", accepted.len() + 1, accept_count);
                        accepted.push(Arc::new(session));
                    }
                    Err(error) => eprintln!("accept error: {error}"),
                }
            }
            accepted
        }
    });

    let client = endpoint;
    let mut sessions: Vec<Arc<QuicSession>> = Vec::new();
    for j in (index + 1)..n {
        // Dial higher-index peers, retrying until their server is up. Each
        // attempt is bounded so a stalled connect (peer not yet listening, or a
        // dropped handshake packet) is retried rather than blocking forever.
        eprintln!("  dialing node {j} at {}", addrs[j]);
        let mut attempt = 0;
        loop {
            let dial = tokio::time::timeout(
                Duration::from_secs(2),
                client.connect(addrs[j], "joule-node"),
            )
            .await;
            match dial {
                Ok(Ok(session)) => {
                    eprintln!("  connected to node {j}");
                    sessions.push(Arc::new(session));
                    break;
                }
                _ if attempt < 100 => {
                    attempt += 1;
                    tokio::time::sleep(Duration::from_millis(150)).await;
                }
                Ok(Err(error)) => return Err(format!("dial node {j}: {error}").into()),
                Err(_) => return Err(format!("dial node {j}: timed out repeatedly").into()),
            }
        }
    }
    eprintln!("  dials done, awaiting {accept_count} accepts");
    let accepted = accept_handle.await?;
    sessions.extend(accepted);
    assert_eq!(sessions.len(), n - 1, "a full mesh has n-1 peers");
    println!("node {index}: mesh formed, {} peers", sessions.len());

    // ── Reader tasks: each peer connection yields exactly two messages. ──
    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<Msg>();
    for session in &sessions {
        let session = session.clone();
        let tx = tx.clone();
        tokio::spawn(async move {
            for _ in 0..2 {
                match recv_msg(&session).await {
                    Ok(msg) => {
                        let _ = tx.send(msg);
                    }
                    Err(error) => {
                        eprintln!("recv error: {error}");
                        break;
                    }
                }
            }
        });
    }
    drop(tx);

    // ── Phase 1: broadcast our proposal, barrier on all peers' proposals. ──
    let (my_op, my_prior) = build_proposal(index);
    let propose = Msg::Propose { from: index, op: my_op.clone(), prior: my_prior.clone() };
    for session in &sessions {
        send_msg(session, &propose).await?;
    }

    let mut proposals: BTreeMap<usize, (SignedOp, ContentId)> = BTreeMap::new();
    proposals.insert(index, (my_op, my_prior));
    let mut stashed_votes: BTreeMap<usize, BTreeSet<String>> = BTreeMap::new();
    while proposals.len() < n {
        match tokio::time::timeout(Duration::from_secs(30), rx.recv()).await {
            Ok(Some(Msg::Propose { from, op, prior })) => {
                proposals.insert(from, (op, prior));
            }
            Ok(Some(Msg::Vote { from, accept })) => {
                stashed_votes.insert(from, accept);
            }
            Ok(None) => return Err("channel closed before all proposals arrived".into()),
            Err(_) => return Err("timed out collecting proposals".into()),
        }
    }
    println!("node {index}: collected {} proposals", proposals.len());

    let proposal_vec: Vec<(SignedOp, ContentId)> = proposals.values().cloned().collect();

    // ── Phase 2: vote on the full set, barrier on all peers' votes. ──
    let my_vote = node.vote(&proposal_vec);
    let vote_msg = Msg::Vote { from: index, accept: my_vote.clone() };
    for session in &sessions {
        send_msg(session, &vote_msg).await?;
    }

    let mut votes: BTreeMap<usize, BTreeSet<String>> = stashed_votes;
    votes.insert(index, my_vote);
    while votes.len() < n {
        match tokio::time::timeout(Duration::from_secs(30), rx.recv()).await {
            Ok(Some(Msg::Vote { from, accept })) => {
                votes.insert(from, accept);
            }
            Ok(Some(Msg::Propose { from, op, prior })) => {
                proposals.entry(from).or_insert((op, prior));
            }
            Ok(None) => return Err("channel closed before all votes arrived".into()),
            Err(_) => return Err("timed out collecting votes".into()),
        }
    }
    println!("node {index}: collected {} vote sets", votes.len());

    // ── Commit deterministically. ──
    let vote_vec: Vec<BTreeSet<String>> = votes.values().cloned().collect();
    node.commit_round(&proposal_vec, &vote_vec);

    println!(
        "node {index}: COMMITTED {} ops, agreed-log hash = {}",
        node.log().len(),
        log_hash(&node),
    );
    for committed in node.log() {
        println!(
            "  seq {} · {}",
            committed.sequence,
            committed.op.op.op_id().expect("op_id").0,
        );
    }

    // Give peers a moment to finish reading before sockets close.
    tokio::time::sleep(Duration::from_millis(500)).await;
    Ok(())
}

// ── Star topology: BFT consensus over QUIC when some hosts cannot dial out ──
//
// On the M4/M5 macOS laptops a process's bound-port socket cannot initiate a
// cross-machine QUIC handshake (a quinn/macOS `sendmsg` EHOSTUNREACH bug), but
// those hosts CAN accept inbound connections and send data on an established
// one, and the Studio can dial out fine. So a send-capable HUB dials every node
// and relays each round's proposals and votes between them.
//
// This does not weaken the safety of the agreement: votes are Ed25519-signed by
// each node's key and every node re-verifies them, so the hub can drop or stall
// messages (a liveness fault, within `f`) but can never forge a vote or change
// the committed set. Each node still validates and commits INDEPENDENTLY, and
// all honest logs stay byte-identical.

#[derive(Serialize, Deserialize)]
enum StarMsg {
    /// A leaf's signed coding-op proposal, sent to the hub.
    Propose { from: usize, op: SignedOp, prior: ContentId },
    /// The hub's broadcast of every node's proposal for the round.
    Proposals { items: Vec<(usize, SignedOp, ContentId)> },
    /// A node's signed accept-vote set, sent to the hub.
    Vote { from: usize, accept: BTreeSet<String>, sig: Vec<u8> },
    /// The hub's broadcast of every node's signed vote.
    Votes { items: Vec<(usize, BTreeSet<String>, Vec<u8>)> },
}

/// Canonical bytes a node signs to authenticate its vote for the round. The
/// op-ids are already sorted (a `BTreeSet`), so this is deterministic.
fn vote_signing_bytes(from: usize, accept: &BTreeSet<String>) -> Vec<u8> {
    let ids: Vec<&str> = accept.iter().map(String::as_str).collect();
    format!("lockstep-vote|{from}|{}", ids.join(",")).into_bytes()
}

fn sign_vote(from: usize, accept: &BTreeSet<String>) -> Vec<u8> {
    use ed25519_dalek::Signer;
    node_key(from).sign(&vote_signing_bytes(from, accept)).to_bytes().to_vec()
}

/// Whether `sig` is a valid signature of `(from, accept)` by node `from`'s key —
/// the check that stops a Byzantine hub from forging or altering a vote.
fn vote_is_authentic(from: usize, accept: &BTreeSet<String>, sig: &[u8], n: usize) -> bool {
    use ed25519_dalek::Verifier;
    if from >= n {
        return false;
    }
    let Ok(signature) = ed25519_dalek::Signature::from_slice(sig) else {
        return false;
    };
    node_key(from)
        .verifying_key()
        .verify(&vote_signing_bytes(from, accept), &signature)
        .is_ok()
}

async fn send_framed<T: Serialize>(
    session: &QuicSession,
    msg: &T,
) -> Result<(), Box<dyn std::error::Error>> {
    let body = serde_json::to_vec(msg)?;
    let mut stream = session.open_uni().await?;
    stream.write_all(&(body.len() as u32).to_be_bytes()).await?;
    stream.write_all(&body).await?;
    stream.finish()?;
    Ok(())
}

async fn recv_framed<T: serde::de::DeserializeOwned>(
    session: &QuicSession,
) -> Result<T, Box<dyn std::error::Error>> {
    let mut stream = session.accept_uni().await?;
    let mut len = [0u8; 4];
    stream.read_exact(&mut len).await?;
    let n = u32::from_be_bytes(len) as usize;
    let mut buf = vec![0u8; n];
    stream.read_exact(&mut buf).await?;
    Ok(serde_json::from_slice(&buf)?)
}

/// Commit the round from the full proposal set and the collected signed votes,
/// counting only authentic ones, then print the agreed-log hash.
fn commit_star(
    node: &mut RoundNode,
    proposals: &[(SignedOp, ContentId)],
    vote_items: &[(usize, BTreeSet<String>, Vec<u8>)],
    n: usize,
) {
    let mut authentic = 0usize;
    let verified: Vec<BTreeSet<String>> = vote_items
        .iter()
        .filter(|(from, accept, sig)| {
            let ok = vote_is_authentic(*from, accept, sig, n);
            if ok {
                authentic += 1;
            }
            ok
        })
        .map(|(_, accept, _)| accept.clone())
        .collect();
    node.commit_round(proposals, &verified);
    println!(
        "star node {}: COMMITTED {} ops from {}/{} authentic votes, agreed-log hash = {}",
        node.index(),
        node.log().len(),
        authentic,
        vote_items.len(),
        log_hash(node),
    );
    for committed in node.log() {
        println!("  seq {} · {}", committed.sequence, committed.op.op.op_id().expect("op_id").0);
    }
}

async fn run_star(args: &[String]) -> Result<(), Box<dyn std::error::Error>> {
    let index: usize = args.get(2).ok_or("need index")?.parse()?;
    let hub: usize = args.get(3).ok_or("need hub index")?.parse()?;
    let addrs: Vec<std::net::SocketAddr> = args
        .get(4)
        .ok_or("need addr list")?
        .split(',')
        .map(|s| s.trim().parse())
        .collect::<Result<_, _>>()?;
    let cert_path = args.get(5).ok_or("need cert path")?;
    let key_path = args.get(6).ok_or("need key path")?;
    let n = addrs.len();
    if index >= n || hub >= n {
        return Err("index/hub out of range".into());
    }
    let cert_der = std::fs::read(cert_path)?;
    let key_der = std::fs::read(key_path)?;
    let membership =
        Membership::new((0..n).map(|i| node_key(i).verifying_key().to_bytes()).collect());
    let mut node = RoundNode::new(index, membership, n);
    println!(
        "star node {index}/{n} (hub={hub}): f={} threshold={}",
        node.f(),
        node.supermajority()
    );

    let listen = std::net::SocketAddr::from(([0, 0, 0, 0], addrs[index].port()));

    if index == hub {
        // Hub: a dual-role endpoint that dials every leaf (the hub runs on a
        // send-capable host) and relays the round.
        let endpoint =
            QuicEndpoint::node_from_der(listen, cert_der.clone(), key_der, cert_der)?;
        let mut leaves: Vec<(usize, Arc<QuicSession>)> = Vec::new();
        for j in 0..n {
            if j == hub {
                continue;
            }
            eprintln!("  hub dialing leaf {j} at {}", addrs[j]);
            let mut attempt = 0;
            loop {
                let dial = tokio::time::timeout(
                    Duration::from_secs(2),
                    endpoint.connect(addrs[j], "joule-node"),
                )
                .await;
                match dial {
                    Ok(Ok(session)) => {
                        leaves.push((j, Arc::new(session)));
                        break;
                    }
                    _ if attempt < 100 => {
                        attempt += 1;
                        tokio::time::sleep(Duration::from_millis(150)).await;
                    }
                    Ok(Err(error)) => return Err(format!("hub dial leaf {j}: {error}").into()),
                    Err(_) => return Err(format!("hub dial leaf {j}: timed out").into()),
                }
            }
            eprintln!("  hub connected to leaf {j}");
        }

        // Phase 1: collect every leaf's proposal (+ the hub's own), broadcast.
        let mut proposals: BTreeMap<usize, (SignedOp, ContentId)> = BTreeMap::new();
        let (op, prior) = build_proposal(hub);
        proposals.insert(hub, (op, prior));
        for (_, session) in &leaves {
            match recv_framed::<StarMsg>(session).await? {
                StarMsg::Propose { from, op, prior } => {
                    proposals.insert(from, (op, prior));
                }
                _ => return Err("hub expected Propose".into()),
            }
        }
        let proposal_items: Vec<(usize, SignedOp, ContentId)> =
            proposals.iter().map(|(i, (o, p))| (*i, o.clone(), p.clone())).collect();
        for (_, session) in &leaves {
            send_framed(session, &StarMsg::Proposals { items: proposal_items.clone() }).await?;
        }

        // Phase 2: collect every node's signed vote, broadcast.
        let proposal_vec: Vec<(SignedOp, ContentId)> = proposals.values().cloned().collect();
        let mut votes: BTreeMap<usize, (BTreeSet<String>, Vec<u8>)> = BTreeMap::new();
        let my_vote = node.vote(&proposal_vec);
        let my_sig = sign_vote(hub, &my_vote);
        votes.insert(hub, (my_vote, my_sig));
        for (_, session) in &leaves {
            match recv_framed::<StarMsg>(session).await? {
                StarMsg::Vote { from, accept, sig } => {
                    votes.insert(from, (accept, sig));
                }
                _ => return Err("hub expected Vote".into()),
            }
        }
        let vote_items: Vec<(usize, BTreeSet<String>, Vec<u8>)> =
            votes.iter().map(|(i, (a, s))| (*i, a.clone(), s.clone())).collect();
        for (_, session) in &leaves {
            send_framed(session, &StarMsg::Votes { items: vote_items.clone() }).await?;
        }

        commit_star(&mut node, &proposal_vec, &vote_items, n);
        tokio::time::sleep(Duration::from_millis(500)).await;
    } else {
        // Leaf: accept the hub's connection and participate over it.
        let endpoint = QuicEndpoint::server_from_der(listen, cert_der, key_der)?;
        eprintln!("  leaf {index} awaiting the hub");
        let session = Arc::new(endpoint.accept().await?);
        eprintln!("  leaf {index} connected to the hub");

        let (op, prior) = build_proposal(index);
        send_framed(&session, &StarMsg::Propose { from: index, op, prior }).await?;
        let proposal_items = match recv_framed::<StarMsg>(&session).await? {
            StarMsg::Proposals { items } => items,
            _ => return Err("leaf expected Proposals".into()),
        };
        let proposal_vec: Vec<(SignedOp, ContentId)> =
            proposal_items.iter().map(|(_, o, p)| (o.clone(), p.clone())).collect();

        let my_vote = node.vote(&proposal_vec);
        let my_sig = sign_vote(index, &my_vote);
        send_framed(&session, &StarMsg::Vote { from: index, accept: my_vote, sig: my_sig }).await?;
        let vote_items = match recv_framed::<StarMsg>(&session).await? {
            StarMsg::Votes { items } => items,
            _ => return Err("leaf expected Votes".into()),
        };

        commit_star(&mut node, &proposal_vec, &vote_items, n);
        tokio::time::sleep(Duration::from_millis(300)).await;
    }
    Ok(())
}