//! Cross-machine MoQT-over-QUIC demo (`draft-ietf-moq-transport-22`, ALPN
//! `moqt-22`): a publisher streams signed, receipt-bearing objects; a subscriber
//! on another host verifies every object against the publisher's pinned Ed25519
//! key over a real QUIC connection. The group's receipts ride its receipts track
//! as one subgroup stream, the objects the media track as another, each object's
//! receipt in its properties.
//!
//! Build:  cargo build --example quic_xnet --features quic
//! Publish (host A):
//!   quic_xnet publish 0.0.0.0:9443 /tmp/jmoq-cert.der /tmp/jmoq-pub.hex
//! Subscribe (host B), after copying the cert + pubkey files over:
//!   quic_xnet subscribe 10.10.10.2:9443 /tmp/jmoq-cert.der <pubkey_hex>
//!
//! The TLS is a pinned self-signed cert (SAN `joule-pub`); the subscriber dials
//! with that server name and trusts only that cert, and trusts only the printed
//! publisher key as a receipt anchor.

#[cfg(not(feature = "quic"))]
fn main() {
    eprintln!("rebuild with --features quic");
    std::process::exit(2);
}

#[cfg(feature = "quic")]
fn main() -> Result<(), Box<dyn std::error::Error>> {
    let rt = tokio::runtime::Builder::new_current_thread().enable_all().build()?;
    rt.block_on(run())
}

#[cfg(feature = "quic")]
fn hex_encode(bytes: &[u8]) -> String {
    bytes.iter().map(|b| format!("{:02x}", b)).collect()
}

#[cfg(feature = "quic")]
fn hex_decode(s: &str) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
    let s = s.trim();
    if s.len() % 2 != 0 {
        return Err("odd-length hex".into());
    }
    (0..s.len())
        .step_by(2)
        .map(|i| u8::from_str_radix(&s[i..i + 2], 16).map_err(|e| e.into()))
        .collect()
}

#[cfg(feature = "quic")]
async fn run() -> Result<(), Box<dyn std::error::Error>> {
    use bytes::Bytes;
    use ed25519_dalek::{SigningKey, VerifyingKey};
    use joule_moq::messages::ObjectHeader;
    use joule_moq::receipt::{ObjectReceipt, ReceiptBuilder, ReceiptVerifier};
    use joule_moq::track::{GroupId, ObjectStatus, TrackAlias};
    use joule_moq::carriage::{ObjectExtras, ReceiptBundle};
    use joule_moq::draft22::{StreamHeader, SubgroupHeader, SubgroupId};
    use joule_moq::transport_quic::{generate_self_signed, send_object, send_receipt_bundle, QuicEndpoint};
    use rand_core_06::OsRng;
    use rustls::pki_types::CertificateDer;

    let args: Vec<String> = std::env::args().collect();
    let mode = args.get(1).map(|s| s.as_str()).unwrap_or("");
    let track_alias = TrackAlias(7);
    let receipts_alias = 8u64;
    let group_id = GroupId(1);
    let origin = "did:key:zJouleMoqPublisher";

    match mode {
        "publish" => {
            let bind: std::net::SocketAddr = args.get(2).ok_or("need bind addr")?.parse()?;
            let cert_out = args.get(3).ok_or("need cert out path")?;
            let pubkey_out = args.get(4).ok_or("need pubkey out path")?;

            let (chain, key) = generate_self_signed(vec!["joule-pub".into()])?;
            std::fs::write(cert_out, chain[0].as_ref())?;
            let signing_key = SigningKey::generate(&mut OsRng);
            let pubkey_hex = hex_encode(&signing_key.verifying_key().to_bytes());
            std::fs::write(pubkey_out, &pubkey_hex)?;

            let server = QuicEndpoint::server(bind, chain, key)?;
            println!("publisher: listening on {bind}");
            println!("publisher: pubkey {pubkey_hex}");
            println!("publisher: cert -> {cert_out}, pubkey -> {pubkey_out}");

            let session = server.accept().await?;
            println!("publisher: subscriber connected from {}", session.remote_address());

            let frames = [
                (Bytes::from_static(b"frame-one"), 100u64),
                (Bytes::from_static(b"frame-two"), 110u64),
                (Bytes::from_static(b"frame-three"), 120u64),
            ];
            let mut builder = ReceiptBuilder::new(track_alias, group_id);
            let mut pending = Vec::new();
            for (payload, joules) in frames.iter() {
                let hash = ObjectReceipt::hash_payload(payload);
                let now = chrono::Utc::now().to_rfc3339();
                let id = builder.add_object(hash, *joules, origin, now);
                pending.push((id, payload.clone(), hash, *joules));
            }
            let proofs: Vec<_> = (0..pending.len()).map(|i| builder.proof_for(i).unwrap()).collect();
            let group_end = chrono::Utc::now().to_rfc3339();
            let group_receipt = builder.finalize(&signing_key, "joule-moq-xnet", group_end);

            let mut sent = Vec::new();
            for ((id, payload, hash, joules), proof) in pending.into_iter().zip(proofs) {
                let header = ObjectHeader {
                    track_alias,
                    group_id,
                    object_id: id,
                    publisher_priority: 128,
                    status: ObjectStatus::Normal,
                    payload_len: payload.len() as u64,
                    mrl_layer: 0,
                    mrl_total_layers: 0,
                    wai_envelope: false,
                };
                let receipt = ObjectReceipt {
                    track_alias,
                    group_id,
                    object_id: id,
                    content_hash: hash,
                    joules_micro: joules,
                    origin: origin.to_string(),
                    merkle_proof: proof,
                    energy_class: None,
                };
                sent.push((header, payload, receipt));
            }
            let bundle = ReceiptBundle { group: group_receipt.clone(), objects: sent.iter().map(|s| s.2.clone()).collect() };
            send_receipt_bundle(&session, receipts_alias, group_id.0, 128, &bundle).await?;
            let mut media = session
                .open_subgroup(SubgroupHeader {
                    track_alias: track_alias.0,
                    group_id: group_id.0,
                    subgroup_id: SubgroupId::Zero,
                    publisher_priority: Some(128),
                    properties: true,
                    end_of_group: true,
                    first_object: true,
                })
                .await?;
            for (header, payload, receipt) in sent {
                send_object(&mut media, &header, &payload, &ObjectExtras { receipt: Some(receipt), wai_receipt: None, delivery_timeout_ms: None }).await?;
            }
            media.finish()?;
            println!("publisher: streamed group receipt + {} objects", group_receipt.leaf_count);
            tokio::time::sleep(std::time::Duration::from_millis(500)).await;
            println!("publisher: done");
            Ok(())
        }
        "subscribe" => {
            let connect: std::net::SocketAddr = args.get(2).ok_or("need connect addr")?.parse()?;
            let cert_in = args.get(3).ok_or("need cert in path")?;
            let pubkey_hex = args.get(4).ok_or("need pubkey hex")?;

            let cert = CertificateDer::from(std::fs::read(cert_in)?);
            let pubkey_bytes: [u8; 32] = hex_decode(pubkey_hex)?
                .try_into()
                .map_err(|_| "pubkey must be 32 bytes")?;
            let publisher_key = VerifyingKey::from_bytes(&pubkey_bytes)?;

            let client = QuicEndpoint::client("0.0.0.0:0".parse()?, vec![cert])?;
            let session = client.connect(connect, "joule-pub").await?;
            println!("subscriber: connected to {connect}");

            let mut receipts = None;
            let mut objects = Vec::new();
            for _ in 0..2 {
                let mut rx = session.accept_data().await?;
                let StreamHeader::Subgroup(h) = rx.header else { return Err("expected a subgroup stream".into()) };
                if h.track_alias == receipts_alias {
                    receipts = Some(rx.read_receipt_bundle().await?);
                } else {
                    while let Some(o) = rx.next_object().await? {
                        objects.push(o);
                    }
                }
            }
            let group_receipt = receipts.ok_or("no receipts track")?.group;

            let verifier = ReceiptVerifier::with_anchors(vec![publisher_key]);
            let root = verifier.verify_group(&group_receipt)?;
            println!(
                "subscriber: group receipt verified vs pinned key — {} objects, {} µJ, root {}",
                group_receipt.leaf_count,
                group_receipt.total_joules_micro,
                hex_encode(&root.root[..8]),
            );

            for (i, (_header, payload, extras)) in objects.into_iter().enumerate() {
                let receipt = extras.receipt.ok_or("object missing its receipt")?;
                verifier.verify_payload(&payload, &receipt, &root)?;
                println!(
                    "subscriber: object {i} VERIFIED — {} ({} µJ)",
                    String::from_utf8_lossy(&payload),
                    receipt.joules_micro,
                );
            }
            println!("subscriber: ALL OBJECTS VERIFIED over real QUIC, cross-machine");
            Ok(())
        }
        _ => {
            eprintln!("usage: quic_xnet publish <bind> <cert_out> <pubkey_out>");
            eprintln!("       quic_xnet subscribe <connect> <cert_in> <pubkey_hex>");
            std::process::exit(2);
        }
    }
}