//! Enhanced RTMP into the room, with a publisher of our own (`tests/common/rtmp.rs`): a
//! stream key from `/rtmp` under a grant, the handshake, `connect` with `capsEx`, and
//! FLV tag bodies built from Enhanced RTMP v2's tables. A subscriber at the relay
//! receives every group, verifies every receipt under the gateway's key (and the relay's
//! door verifies them too), and gets each frame in the payload framing
//! `live-ingest.md` §12's table gives it, at the room-clock time §7 maps it to.
//!
//! - legacy AVC and AAC: Annex B with the record's SPS and PPS before each IDR access
//!   unit, one ADTS frame per raw AAC frame;
//! - enhanced HEVC (CodedFrames with a composition time offset, `live.cto`) and Opus;
//! - AV1 and FLAC, VP9 and enhanced AAC, a ModEx nanosecond offset;
//! - multitrack: under a grant whose video profile has two layers, each track carried as
//!   the layer its `onMetaData` entry matches, and the tracks that match none, or are not
//!   audio track 0, counted and not carried;
//! - RTMPS; the refusals of a stream key; the reconnect request at shutdown.

mod common;

use common::rtmp::*;
use common::*;
// Both modules have an `adts`: this file's frames are the RTMP helper's.
use common::rtmp::adts;
use joule_moq_ingest::frame::{self, Codec};
use joule_moq_ingest::room::media;
use joule_moq_ingest::rtmp::amf0::Value;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::time::timeout;

const GOP: u64 = 30;
const FRAMES: u64 = 90;

/// A world with RTMP listening and a skew bound wide enough that a publisher sending
/// faster than real time is mapped by its timestamps alone (§7).
async fn world(f: impl FnOnce(&mut joule_moq_ingest::GatewayConfig)) -> World {
    World::start(|c| {
        c.rtmp.listen = Some("127.0.0.1:0".parse().unwrap());
        c.timeline.max_skew_ms = 600_000;
        f(c)
    })
    .await
}

async fn client(w: &World) -> Client<TcpStream> {
    let s = TcpStream::connect(w.gateway.rtmp_addr.expect("an RTMP listener")).await.unwrap();
    Client::start(s).await
}

/// Publish `key`, and fail the test with the gateway's description where it refuses.
async fn publishing(w: &World, key: &str, caps_ex: u32) -> Client<TcpStream> {
    let mut c = client(w).await;
    let r = c.connect(caps_ex).await;
    let props = &r[2];
    assert_eq!(props.get("capsEx").and_then(Value::as_f64), Some(15.0), "Reconnect, Multitrack, ModEx, TimestampNanoOffset");
    assert!(props.get("videoFourCcInfoMap").and_then(|m| m.get("hvc1")).is_some());
    let s = c.publish(key).await;
    assert_eq!((s.code.as_str(), s.level.as_str()), ("NetStream.Publish.Start", "status"), "{}", s.description);
    c
}

/// End the publish and wait for every track to end at the relay.
async fn finish(w: &World, c: &mut Client<TcpStream>, watcher: &Watcher, n: usize) {
    c.delete_stream().await;
    assert!(until(STEP, || watcher.ended.lock().unwrap().len() == n).await, "PUBLISH_DONE on every track");
    settle(watcher, Duration::from_millis(400)).await;
    c.shutdown().await;
    assert!(until(STEP, || !w.gateway.gateway.records().is_empty()).await, "the session's record");
}

fn track(wt: &Watcher, kind: u8, layer: u8) -> &Watched {
    wt.tracks.iter().find(|t| t.kind == kind && t.layer == layer).unwrap()
}

/// The room-clock time of each object of `t`, and its `live` member.
fn lives(wt: &Watcher, t: &Watched) -> Vec<serde_json::Value> {
    envelopes(wt, t).into_iter().map(|e| e.manifest["live"].clone()).collect()
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn legacy_avc_and_aac() {
    timeout(Duration::from_secs(120), legacy_avc_and_aac_at(22)).await.expect("in time");
}

/// The same publish with the gateway's sessions to the relay at -23, -18 and -16: the
/// relay bridges each to its -22 model, the watcher reads at -22, every object verifies at
/// the door and at the watcher, and the relay's summary of the gateway's session names the
/// draft it spoke.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn legacy_avc_and_aac_at_23_18_and_16() {
    for d in [23, 18, 16] {
        timeout(Duration::from_secs(120), legacy_avc_and_aac_at(d)).await.unwrap_or_else(|_| panic!("draft {d} in time"));
    }
}

async fn legacy_avc_and_aac_at(draft: u32) {
    {
        let w = world(|c| c.relay.draft = draft).await;
        let auth = w.auth(&w.publisher, 0);
        let key = stream_key(&w, &auth).await;
        let mut c = publishing(&w, &key, 0).await;
        let watcher = watch(&w).await;
        c.metadata(vec![("width".into(), Value::Number(640.0)), ("height".into(), Value::Number(360.0)), ("framerate".into(), Value::Number(30.0)), ("videocodecid".into(), Value::Number(7.0)), ("audiocodecid".into(), Value::Number(10.0))]).await;
        c.video(0, legacy_avc(0, true, 0, &avc_record())).await;
        c.audio(0, legacy_aac(0, &AAC_ASC)).await;
        let (mut want_v, mut want_a) = (Vec::new(), Vec::new());
        for i in 0..FRAMES {
            let key = i % GOP == 0;
            let nals = avc_nals(i, key);
            c.video((i * 1000 / 30) as u32, legacy_avc(1, key, 0, &lengths(&nals))).await;
            let mut au = if key { vec![AVC_SPS.to_vec(), AVC_PPS.to_vec()] } else { Vec::new() };
            au.extend(nals);
            want_v.push(annexb(&au));
            let raw = aac_raw(i);
            c.audio((i * 1000 / 30) as u32, legacy_aac(1, &raw)).await;
            want_a.push(adts(&raw));
        }
        finish(&w, &mut c, &watcher, 2).await;
        let (vf, vg) = verified(&watcher, track(&watcher, media::VIDEO, 0), &w.gateway_key, Codec::Avc);
        let (af, _) = verified(&watcher, track(&watcher, media::AUDIO, 0), &w.gateway_key, Codec::Aac);
        same(&vf, &want_v, "AVC: Annex B, the record's parameter sets before each IDR");
        same(&af, &want_a, "AAC: one ADTS frame each");
        assert_eq!(vg as u64, FRAMES / GOP, "one video group per keyframe");
        // §7: t follows the message timestamps from the first frame's arrival.
        let l = lives(&watcher, track(&watcher, media::VIDEO, 0));
        let t0 = l[0]["t"].as_u64().unwrap();
        for (i, x) in l.iter().enumerate() {
            assert_eq!(x["t"].as_u64().unwrap() - t0, (i as u64 * 1000 / 30) * 1000, "frame {i}'s room-clock time");
            assert_eq!(x["rate"], 1_000_000);
            assert!(x.get("cto").is_none(), "no composition offset when it is 0");
        }
        assert_eq!(w.relay.relay().stats().refused_at_door, 0, "the relay's door verified every object");
        let rec = w.gateway.gateway.records().remove(0);
        rec.verify().expect("the record verifies under the gateway's key");
        assert_eq!((rec.record.protocol.as_str(), rec.record.role, rec.record.ended.as_str()), ("ertmp", 3, "deleted"));
        let out: u64 = rec.record.tracks.iter().filter(|t| t.dir == 2).map(|t| t.objects).sum();
        assert_eq!(out, 2 * FRAMES);
        let name = format!("moqt-{draft}");
        assert!(until(STEP, || w.relay.relay().summaries().iter().any(|x| x.summary.draft == name)).await, "a session at {name}");
        w.stop().await;
    }
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn enhanced_hevc_with_composition_offsets_and_opus() {
    timeout(Duration::from_secs(120), async {
        let w = world(|_| {}).await;
        let auth = w.auth(&w.publisher, 0);
        let key = stream_key(&w, &auth).await;
        let mut c = publishing(&w, &key, 0x0f).await;
        let watcher = watch(&w).await;
        c.metadata(vec![("width".into(), Value::Number(1280.0)), ("height".into(), Value::Number(720.0)), ("framerate".into(), Value::Number(30.0)), ("videocodecid".into(), Value::Number(u32::from_be_bytes(*b"hvc1") as f64))]).await;
        c.video(0, ex_video(true, 0, b"hvc1", &hevc_record())).await;
        c.audio(0, ex_audio(0, b"Opus", &opus_head())).await;
        let (mut want_v, mut want_a) = (Vec::new(), Vec::new());
        for i in 0..FRAMES {
            let key = i % GOP == 0;
            let nals = hevc_nals(i, key);
            let ts = (i * 1000 / 30) as u32;
            // CodedFrames with a composition time offset of 66 ms, or CodedFramesX.
            if i % 2 == 0 {
                let mut b = 66i32.to_be_bytes()[1..].to_vec();
                b.extend(lengths(&nals));
                c.video(ts, ex_video(key, 1, b"hvc1", &b)).await;
            } else {
                c.video(ts, ex_video(key, 3, b"hvc1", &lengths(&nals))).await;
            }
            let mut au = if key { hevc_sets() } else { Vec::new() };
            au.extend(nals);
            want_v.push(annexb(&au));
            let p = opus_packet(i);
            c.audio(ts, ex_audio(1, b"Opus", &p)).await;
            want_a.push(p);
        }
        finish(&w, &mut c, &watcher, 2).await;
        let (vf, vg) = verified(&watcher, track(&watcher, media::VIDEO, 0), &w.gateway_key, Codec::Hevc);
        let (af, _) = verified(&watcher, track(&watcher, media::AUDIO, 0), &w.gateway_key, Codec::Opus);
        same(&vf, &want_v, "HEVC: Annex B, VPS, SPS and PPS before each IRAP access unit");
        same(&af, &want_a, "Opus packets");
        assert_eq!(vg as u64, FRAMES / GOP);
        let l = lives(&watcher, track(&watcher, media::VIDEO, 0));
        assert_eq!(l[0]["cto"], 66_000, "CodedFrames' offset, in µs");
        assert!(l[1].get("cto").is_none(), "CodedFramesX: an offset of 0, not written");
        let env = envelopes(&watcher, track(&watcher, media::VIDEO, 0));
        assert_eq!(env[0].manifest["model_requirement"]["capability"], "wai.video.hevc");
        assert!(frame::hevc_is_irap(&env[0].payload));
        w.stop().await;
    })
    .await
    .expect("in time");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn av1_with_flac_and_vp9_with_aac_and_a_nano_offset() {
    timeout(Duration::from_secs(150), async {
        for (vc, ac) in [(Codec::Av1, Codec::Flac), (Codec::Vp9, Codec::Aac)] {
            let w = world(|_| {}).await;
            let auth = w.auth(&w.publisher, 0);
            let key = stream_key(&w, &auth).await;
            let mut c = publishing(&w, &key, 0x0f).await;
            let watcher = watch(&w).await;
            let (vf4, af4) = match (vc, ac) {
                (Codec::Av1, _) => (b"av01", b"fLaC"),
                _ => (b"vp09", b"mp4a"),
            };
            c.metadata(vec![("width".into(), Value::Number(640.0)), ("height".into(), Value::Number(360.0)), ("framerate".into(), Value::Number(30.0))]).await;
            let rec = if vc == Codec::Av1 { av1_record(640, 360) } else { vp9_record() };
            c.video(0, ex_video(true, 0, vf4, &rec)).await;
            let aseq = if ac == Codec::Flac { flac_header() } else { AAC_ASC.to_vec() };
            c.audio(0, ex_audio(0, af4, &aseq)).await;
            let mut want_v = Vec::new();
            let mut want_a = Vec::new();
            for i in 0..FRAMES {
                let key = i % GOP == 0;
                let ts = (i * 1000 / 30) as u32;
                let f = if vc == Codec::Av1 { av1_tu(i, key) } else { vp9_frame(i, key, 640, 360) };
                // Every tenth frame carries a ModEx nanosecond offset of 400 µs.
                if i % 10 == 5 {
                    c.video(ts, ex_video_nano(key, 1, vf4, 400_000, &f)).await;
                } else {
                    c.video(ts, ex_video(key, 1, vf4, &f)).await;
                }
                want_v.push(if vc == Codec::Av1 {
                    let mut o = if key { av1_sequence_header(640, 360) } else { Vec::new() };
                    o.extend_from_slice(&f[2..]);
                    o
                } else {
                    f
                });
                let (a, want) = if ac == Codec::Flac {
                    let fr = flac_frame(i);
                    let mut s = flac_header();
                    s[4] |= 0x80;
                    s.extend_from_slice(&fr);
                    (fr, s)
                } else {
                    let raw = aac_raw(i);
                    (raw.clone(), adts(&raw))
                };
                c.audio(ts, ex_audio(1, af4, &a)).await;
                want_a.push(want);
            }
            finish(&w, &mut c, &watcher, 2).await;
            let (vf, vg) = verified(&watcher, track(&watcher, media::VIDEO, 0), &w.gateway_key, vc);
            let (af, _) = verified(&watcher, track(&watcher, media::AUDIO, 0), &w.gateway_key, ac);
            same(&vf, &want_v, &format!("{vc:?}"));
            same(&af, &want_a, &format!("{ac:?}"));
            assert_eq!(vg as u64, FRAMES / GOP);
            let env = envelopes(&watcher, track(&watcher, media::VIDEO, 0));
            let size = match vc {
                Codec::Vp9 => (u16::from_le_bytes([env[0].payload[12], env[0].payload[13]]) as u32, u16::from_le_bytes([env[0].payload[14], env[0].payload[15]]) as u32),
                _ => (u32::from_le_bytes(env[0].payload[8..12].try_into().unwrap()), u32::from_le_bytes(env[0].payload[4..8].try_into().unwrap())),
            };
            assert_eq!(size, (640, 360), "the framing states the size the bitstream gives");
            let l = lives(&watcher, track(&watcher, media::VIDEO, 0));
            assert_eq!(l[5]["rtp"].as_u64().unwrap() - l[4]["rtp"].as_u64().unwrap(), 33_000 + 400, "the nanosecond offset, in µs, on frame 5's media time");
            w.stop().await;
        }
    })
    .await
    .expect("in time");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn multitrack_layers_by_the_grant_s_profile() {
    timeout(Duration::from_secs(120), async {
        let w = world(|_| {}).await;
        // A grant with the room's `camera-720` profile for video: layer 0 640×360 at 30,
        // layer 1 1280×720 at 30, codecs av01 and avc1.
        let mut g = w.grant(&w.publisher, 0);
        g.profiles = vec![1];
        let g = g.sign(&w.authority);
        let auth = Auth::new(g, &w.publisher);
        let key = stream_key(&w, &auth).await;
        let mut c = publishing(&w, &key, 0x0f).await;
        let watcher = watch_layers(&w, &[(media::VIDEO, 0), (media::VIDEO, 1), (media::AUDIO, 0)]).await;
        let info = |w: f64, h: f64| Value::obj(&[("width", Value::Number(w)), ("height", Value::Number(h)), ("framerate", Value::Number(30.0))]);
        c.metadata(vec![
            ("width".into(), Value::Number(1280.0)),
            ("height".into(), Value::Number(720.0)),
            ("framerate".into(), Value::Number(29.97)),
            ("videoTrackIdInfoMap".into(), Value::obj(&[("1", info(640.0, 360.0)), ("2", info(1920.0, 1080.0))])),
        ])
        .await;
        c.video(0, ex_video_tracks(true, 0, &[(b"avc1", 0, avc_record()), (b"av01", 1, av1_record(640, 360)), (b"avc1", 2, avc_record())])).await;
        c.audio(0, ex_audio_tracks(0, &[(b"mp4a", 0, AAC_ASC.to_vec()), (b"Opus", 1, opus_head())])).await;
        let (mut hi, mut lo) = (Vec::new(), Vec::new());
        for i in 0..60u64 {
            let key = i % GOP == 0;
            let ts = (i * 1000 / 30) as u32;
            let nals = avc_nals(i, key);
            let mut avc_body = vec![0, 0, 0];
            avc_body.extend(lengths(&nals));
            let tu = av1_tu(i, key);
            c.video(ts, ex_video_tracks(key, 1, &[(b"avc1", 0, avc_body.clone()), (b"av01", 1, tu.clone()), (b"avc1", 2, avc_body)])).await;
            let mut au = if key { vec![AVC_SPS.to_vec(), AVC_PPS.to_vec()] } else { Vec::new() };
            au.extend(nals);
            hi.push(annexb(&au));
            let mut o = if key { av1_sequence_header(640, 360) } else { Vec::new() };
            o.extend_from_slice(&tu[2..]);
            lo.push(o);
            c.audio(ts, ex_audio_tracks(1, &[(b"mp4a", 0, aac_raw(i)), (b"Opus", 1, opus_packet(i))])).await;
        }
        finish(&w, &mut c, &watcher, 3).await;
        let (l0, _) = verified(&watcher, track(&watcher, media::VIDEO, 0), &w.gateway_key, Codec::Av1);
        let (l1, _) = verified(&watcher, track(&watcher, media::VIDEO, 1), &w.gateway_key, Codec::Avc);
        let (a0, _) = verified(&watcher, track(&watcher, media::AUDIO, 0), &w.gateway_key, Codec::Aac);
        same(&l0, &lo, "layer 0: track 1, 640×360 AV1");
        same(&l1, &hi, "layer 1: track 0, 1280×720 at 29.97 (30 rounded) AVC");
        assert_eq!(a0.len(), 60, "audio track 0");
        let rec = w.gateway.gateway.records().remove(0);
        let counts = &rec.record.detail["counts"];
        assert_eq!(counts["track_no_layer"], 61, "video track 2 (1920×1080, no layer of the profile), every packet");
        assert_eq!(counts["audio_track_not_default"], 61, "audio track 1, every packet");
        w.stop().await;
    })
    .await
    .expect("in time");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn stream_key_refusals_and_admission() {
    timeout(Duration::from_secs(120), async {
        let w = world(|c| c.rtmp.stream_key_ttl_s = 1).await;
        let auth = w.auth(&w.publisher, 0);
        // A viewer's grant publishes nothing: 403.
        let viewer = ed25519_dalek::SigningKey::from_bytes(&[0x61; 32]);
        let va = w.auth(&viewer, 1);
        let r = va.send("POST", &format!("{}/rtmp", w.base()), &[], b"").await;
        assert_eq!((r.status, r.header("www-authenticate").is_some_and(|h| h.contains("insufficient_scope"))), (403, true), "{}", r.text());
        // Without a proof of possession: 401 pop_missing.
        let r = auth.send_raw("POST", &format!("{}/rtmp", w.base()), &[], b"", None).await;
        assert!(r.status == 401 && r.text().contains("pop_missing"), "{} {}", r.status, r.text());
        // A key the gateway did not issue.
        let mut c = client(&w).await;
        c.connect(0).await;
        let s = c.publish("00112233445566778899aabbccddeeff").await;
        assert_eq!((s.code.as_str(), s.level.as_str(), s.description.as_str()), ("NetStream.Publish.BadName", "error", "stream_key_unknown"));
        assert!(c.next().await.is_none(), "the connection is closed");
        // A key used once is not good twice.
        let key = stream_key(&w, &auth).await;
        let mut c = publishing(&w, &key, 0).await;
        let mut c2 = client(&w).await;
        c2.connect(0).await;
        assert_eq!(c2.publish(&key).await.description, "stream_key_unknown");
        c.delete_stream().await;
        c.drain().await;
        // A key past its time.
        let key = stream_key(&w, &auth).await;
        tokio::time::sleep(Duration::from_millis(1_300)).await;
        let mut c = client(&w).await;
        c.connect(0).await;
        assert_eq!(c.publish(&key).await.description, "stream_key_expired");
        w.stop().await;
        // A grant whose window has closed by the time its key is used is refused for it
        // (§3.2 step 3, checked again at publish).
        let w = world(|_| {}).await;
        let s = now_us() / 1_000_000;
        let short = joule_moq_ingest::room::build::grant(&w.descriptor, &w.authority, &w.publisher.verifying_key().to_bytes(), 0, s - 60, s + 2);
        let key = stream_key(&w, &Auth::new(short, &w.publisher)).await;
        tokio::time::sleep(Duration::from_millis(2_600)).await;
        let mut c = client(&w).await;
        c.connect(0).await;
        assert_eq!(c.publish(&key).await.description, "outside_grant_window");
        w.stop().await;
    })
    .await
    .expect("in time");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn rtmps_and_the_reconnect_request() {
    timeout(Duration::from_secs(120), async {
        let id = joule_moq_relay::cert::RelayIdentity::dev_self_signed(&["localhost".to_string()]).unwrap();
        let chain: Vec<Vec<u8>> = id.chain.iter().map(|c| c.to_vec()).collect();
        let key_der = id.key.secret_der().to_vec();
        let root = chain[0].clone();
        let w = world(move |c| {
            c.rtmp.tls_listen = Some("127.0.0.1:0".parse().unwrap());
            c.rtmp.tls_der = Some((chain, key_der));
            c.rtmp.reconnect_url = "rtmps://ingest-2.example/live".into();
            c.rtmp.reconnect_grace_ms = 2_000;
        })
        .await;
        let auth = w.auth(&w.publisher, 0);
        let key = stream_key(&w, &auth).await;
        let mut roots = rustls::RootCertStore::empty();
        roots.add(rustls::pki_types::CertificateDer::from(root)).unwrap();
        let cfg = rustls::ClientConfig::builder_with_provider(std::sync::Arc::new(rustls::crypto::ring::default_provider())).with_safe_default_protocol_versions().unwrap().with_root_certificates(roots).with_no_client_auth();
        let tcp = TcpStream::connect(w.gateway.rtmps_addr.unwrap()).await.unwrap();
        let tls = tokio_rustls::TlsConnector::from(std::sync::Arc::new(cfg)).connect(rustls::pki_types::ServerName::try_from("localhost").unwrap(), tcp).await.expect("the TLS handshake");
        let mut c = Client::start(tls).await;
        c.connect(0x01).await;
        let s = c.publish(&key).await;
        assert_eq!(s.code, "NetStream.Publish.Start");
        let watcher = watch(&w).await;
        c.video(0, legacy_avc(0, true, 0, &avc_record())).await;
        for i in 0..30u64 {
            c.video((i * 33) as u32, legacy_avc(1, i == 0, 0, &lengths(&avc_nals(i, i == 0)))).await;
        }
        assert!(until(STEP, || w.gateway.gateway.rtmp_frames().iter().any(|(_, n)| *n >= 30)).await, "every frame reached the gateway");
        // Shut the gateway down: the publisher declared Reconnect, so it is asked to.
        let gw = w.gateway.gateway.clone();
        let gateway = w.gateway;
        let stop = tokio::spawn(async move { gateway.shutdown().await });
        let s = c.status().await.expect("a status");
        assert_eq!((s.code.as_str(), s.level.as_str()), ("NetConnection.Connect.ReconnectRequest", "status"));
        let tc = c.received.iter().rev().find_map(|m| {
            let v = joule_moq_ingest::rtmp::amf0::read_all(&m.payload).ok()?;
            v.get(3)?.get("tcUrl")?.as_str().map(str::to_string)
        });
        assert_eq!(tc.as_deref(), Some("rtmps://ingest-2.example/live"));
        c.shutdown().await;
        stop.await.unwrap();
        assert!(until(STEP, || watcher.ended.lock().unwrap().len() == 2).await, "the tracks end");
        let (vf, _) = verified(&watcher, &watcher.tracks[1], &w.gateway_key, Codec::Avc);
        assert_eq!(vf.len(), 30);
        assert!(gw.stream_keys.lock().unwrap().contains_key(&key), "the key is good once more, for the reconnection");
        let rec = gw.records().remove(0);
        assert_eq!(rec.record.ended, "reconnect_requested");
        assert_eq!(rec.record.detail["tls"], true);
        w.relay.shutdown().await;
    })
    .await
    .expect("in time");
}