//! An Enhanced RTMP publisher of our own, for the gateway's tests: the handshake, the
//! commands, `onMetaData`, and FLV tag bodies, legacy and enhanced, built from Enhanced
//! RTMP v2's tables; and the synthetic frames and configuration records it sends.

#![allow(dead_code)]

use super::{Auth, World, STEP};
use joule_moq_ingest::rtmp::amf0::{self, Value};
use joule_moq_ingest::rtmp::chunk::{self, types, Decoder, Message};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::time::timeout;

/// A stream key from the gateway's `/rtmp` resource, under `auth`'s grant.
pub async fn stream_key(w: &World, auth: &Auth) -> String {
    let r = auth.send("POST", &format!("{}/rtmp", w.base()), &[], b"").await;
    assert_eq!(r.status, 201, "POST /rtmp: {}", r.text());
    let v: serde_json::Value = serde_json::from_slice(&r.body).unwrap();
    v["stream_key"].as_str().unwrap().to_string()
}

/// One status the gateway sent: its code, level and description.
#[derive(Clone, Debug)]
pub struct Status {
    pub code: String,
    pub level: String,
    pub description: String,
}

pub struct Client<S> {
    io: S,
    dec: Decoder,
    out_chunk: usize,
    pub received: Vec<Message>,
}

impl<S: AsyncRead + AsyncWrite + Unpin> Client<S> {
    /// The simple handshake over `io`.
    pub async fn start(mut io: S) -> Client<S> {
        let mut c0c1 = vec![3u8];
        c0c1.extend_from_slice(&[0; 8]);
        c0c1.extend((0..chunk::HANDSHAKE_LEN - 8).map(|i| (i * 13) as u8));
        io.write_all(&c0c1).await.unwrap();
        let mut s = vec![0u8; 1 + 2 * chunk::HANDSHAKE_LEN];
        timeout(STEP, io.read_exact(&mut s)).await.expect("S0 S1 S2 in time").unwrap();
        assert_eq!(s[0], 3);
        assert_eq!(&s[1 + chunk::HANDSHAKE_LEN + 8..], &c0c1[9..], "S2 echoes C1");
        io.write_all(&s[1..1 + chunk::HANDSHAKE_LEN]).await.unwrap();
        Client { io, dec: Decoder::new(1 << 24, 1 << 24), out_chunk: 128, received: Vec::new() }
    }

    pub async fn send(&mut self, m: Message) {
        let b = chunk::encode(&m, self.out_chunk);
        self.io.write_all(&b).await.unwrap();
        if m.type_id == types::SET_CHUNK_SIZE {
            self.out_chunk = u32::from_be_bytes(m.payload[..4].try_into().unwrap()) as usize;
        }
    }

    pub async fn command(&mut self, msid: u32, values: &[Value]) {
        self.send(Message { csid: 3, type_id: types::COMMAND_AMF0, msid, ts: 0, payload: amf0::encode(values) }).await;
    }

    /// The next message the gateway sends, bounded.
    pub async fn next(&mut self) -> Option<Message> {
        let mut buf = vec![0u8; 8192];
        loop {
            if let Some(m) = self.dec.next_message().expect("the gateway's chunk stream reads") {
                if m.type_id == types::SET_CHUNK_SIZE {
                    self.dec.set_chunk_size(u32::from_be_bytes(m.payload[..4].try_into().unwrap())).unwrap();
                }
                self.received.push(m.clone());
                return Some(m);
            }
            match timeout(STEP, self.io.read(&mut buf)).await {
                Ok(Ok(0)) | Ok(Err(_)) | Err(_) => return None,
                Ok(Ok(n)) => self.dec.feed(&buf[..n]),
            }
        }
    }

    /// The values of the next command message, skipping the others.
    pub async fn next_command(&mut self) -> Option<Vec<Value>> {
        while let Some(m) = self.next().await {
            if m.type_id == types::COMMAND_AMF0 {
                return Some(amf0::read_all(&m.payload).expect("a command reads"));
            }
        }
        None
    }

    /// The next `onStatus`.
    pub async fn status(&mut self) -> Option<Status> {
        while let Some(v) = self.next_command().await {
            if v.first().and_then(Value::as_str) == Some("onStatus") {
                let info = v.get(3)?;
                let s = |k| info.get(k).and_then(Value::as_str).unwrap_or("").to_string();
                return Some(Status { code: s("code"), level: s("level"), description: s("description") });
            }
        }
        None
    }

    /// `connect` with `capsEx` and the FourCCs it can send, then the `_result`, and a
    /// chunk size of 4096 from here.
    pub async fn connect(&mut self, caps_ex: u32) -> Vec<Value> {
        let obj = Value::obj(&[
            ("app", Value::str("live")),
            ("type", Value::str("nonprivate")),
            ("tcUrl", Value::str("rtmp://127.0.0.1/live")),
            ("capsEx", Value::Number(caps_ex as f64)),
            ("fourCcList", Value::StrictArray(["av01", "vp09", "vp08", "hvc1", "avc1", "Opus", "fLaC", "mp4a"].iter().map(|f| Value::str(f)).collect())),
        ]);
        self.command(0, &[Value::str("connect"), Value::Number(1.0), obj]).await;
        let r = loop {
            let v = self.next_command().await.expect("the connect _result");
            if v.first().and_then(Value::as_str) == Some("_result") {
                break v;
            }
        };
        self.send(chunk::control(types::SET_CHUNK_SIZE, 4096u32.to_be_bytes().to_vec())).await;
        r
    }

    /// `createStream` and `publish` with `key`; the gateway's status.
    pub async fn publish(&mut self, key: &str) -> Status {
        self.command(0, &[Value::str("releaseStream"), Value::Number(2.0), Value::Null, Value::str(key)]).await;
        self.command(0, &[Value::str("FCPublish"), Value::Number(3.0), Value::Null, Value::str(key)]).await;
        self.command(0, &[Value::str("createStream"), Value::Number(4.0), Value::Null]).await;
        let id = loop {
            let v = self.next_command().await.expect("the createStream _result");
            if v.first().and_then(Value::as_str) == Some("_result") && v.get(1).and_then(Value::as_f64) == Some(4.0) {
                break v.get(3).and_then(Value::as_f64).unwrap() as u32;
            }
        };
        self.command(id, &[Value::str("publish"), Value::Number(5.0), Value::Null, Value::str(key), Value::str("live")]).await;
        self.status().await.expect("a status for publish")
    }

    /// `@setDataFrame` `onMetaData` with `pairs`.
    pub async fn metadata(&mut self, pairs: Vec<(String, Value)>) {
        let p = amf0::encode(&[Value::str("@setDataFrame"), Value::str("onMetaData"), Value::EcmaArray(pairs)]);
        self.send(Message { csid: 4, type_id: types::DATA_AMF0, msid: 1, ts: 0, payload: p }).await;
    }

    pub async fn video(&mut self, ts: u32, body: Vec<u8>) {
        self.send(Message { csid: 6, type_id: types::VIDEO, msid: 1, ts, payload: body }).await;
    }

    pub async fn audio(&mut self, ts: u32, body: Vec<u8>) {
        self.send(Message { csid: 4, type_id: types::AUDIO, msid: 1, ts, payload: body }).await;
    }

    pub async fn delete_stream(&mut self) {
        self.command(0, &[Value::str("FCUnpublish"), Value::Number(6.0), Value::Null, Value::str("")]).await;
        self.command(0, &[Value::str("deleteStream"), Value::Number(7.0), Value::Null, Value::Number(1.0)]).await;
    }

    pub async fn shutdown(&mut self) {
        let _ = self.io.shutdown().await;
    }

    /// Read until the gateway closes the connection, bounded.
    pub async fn drain(&mut self) {
        while self.next().await.is_some() {}
    }
}

// ── FLV tag bodies (Enhanced RTMP v2, Enhanced Video and Enhanced Audio) ──────────

pub fn legacy_avc(packet_type: u8, key: bool, cts: i32, body: &[u8]) -> Vec<u8> {
    let mut o = vec![if key { 0x17 } else { 0x27 }, packet_type];
    o.extend_from_slice(&cts.to_be_bytes()[1..]);
    o.extend_from_slice(body);
    o
}

pub fn legacy_aac(packet_type: u8, body: &[u8]) -> Vec<u8> {
    let mut o = vec![0xaf, packet_type];
    o.extend_from_slice(body);
    o
}

/// An enhanced video body: `isExVideoHeader`, the frame type, the packet type, the
/// FourCC, then `body` (which, for CodedFrames of avc1, hvc1 and vvc1, begins with its
/// SI24 composition time offset).
pub fn ex_video(key: bool, packet_type: u8, fourcc: &[u8; 4], body: &[u8]) -> Vec<u8> {
    let mut o = vec![0x80 | if key { 0x10 } else { 0x20 } | packet_type];
    o.extend_from_slice(fourcc);
    o.extend_from_slice(body);
    o
}

/// The same, after a ModEx `TimestampOffsetNano` of `nano`.
pub fn ex_video_nano(key: bool, packet_type: u8, fourcc: &[u8; 4], nano: u32, body: &[u8]) -> Vec<u8> {
    let mut o = vec![0x80 | if key { 0x10 } else { 0x20 } | 7, 2];
    o.extend_from_slice(&nano.to_be_bytes()[1..]);
    o.push(packet_type);
    o.extend_from_slice(fourcc);
    o.extend_from_slice(body);
    o
}

/// A multitrack video body of type ManyTracksManyCodecs: each track its FourCC, its id,
/// its size and its body.
pub fn ex_video_tracks(key: bool, packet_type: u8, tracks: &[(&[u8; 4], u8, Vec<u8>)]) -> Vec<u8> {
    let mut o = vec![0x80 | if key { 0x10 } else { 0x20 } | 6, 2 << 4 | packet_type];
    for (f, id, b) in tracks {
        o.extend_from_slice(*f);
        o.push(*id);
        o.extend_from_slice(&(b.len() as u32).to_be_bytes()[1..]);
        o.extend_from_slice(b);
    }
    o
}

pub fn ex_audio(packet_type: u8, fourcc: &[u8; 4], body: &[u8]) -> Vec<u8> {
    let mut o = vec![0x90 | packet_type];
    o.extend_from_slice(fourcc);
    o.extend_from_slice(body);
    o
}

/// A multitrack audio body of type ManyTracksManyCodecs.
pub fn ex_audio_tracks(packet_type: u8, tracks: &[(&[u8; 4], u8, Vec<u8>)]) -> Vec<u8> {
    let mut o = vec![0x90 | 5, 2 << 4 | packet_type];
    for (f, id, b) in tracks {
        o.extend_from_slice(*f);
        o.push(*id);
        o.extend_from_slice(&(b.len() as u32).to_be_bytes()[1..]);
        o.extend_from_slice(b);
    }
    o
}

// ── Synthetic configurations and frames ────────────────────────────────────────────

/// Bytes none of which is zero, so no start code appears inside a NAL unit.
pub fn body(seed: u64, n: usize) -> Vec<u8> {
    (0..n).map(|j| ((seed.wrapping_mul(31).wrapping_add(j as u64 * 7)) % 255 + 1) as u8).collect()
}

/// NAL units, each after a four-byte length.
pub fn lengths(nals: &[Vec<u8>]) -> Vec<u8> {
    let mut o = Vec::new();
    for n in nals {
        o.extend_from_slice(&(n.len() as u32).to_be_bytes());
        o.extend_from_slice(n);
    }
    o
}

/// NAL units as an Annex B access unit, each after a four-byte start code.
pub fn annexb(nals: &[Vec<u8>]) -> Vec<u8> {
    let mut o = Vec::new();
    for n in nals {
        o.extend_from_slice(&[0, 0, 0, 1]);
        o.extend_from_slice(n);
    }
    o
}

pub const AVC_SPS: [u8; 6] = [0x67, 0x64, 0x00, 0x1f, 0xac, 0xd9];
pub const AVC_PPS: [u8; 4] = [0x68, 0xeb, 0xe3, 0xcb];

/// An AVCDecoderConfigurationRecord: High profile, level 3.1, four-byte lengths.
pub fn avc_record() -> Vec<u8> {
    let mut o = vec![1, 0x64, 0x00, 0x1f, 0xff, 0xe1];
    o.extend_from_slice(&(AVC_SPS.len() as u16).to_be_bytes());
    o.extend_from_slice(&AVC_SPS);
    o.push(1);
    o.extend_from_slice(&(AVC_PPS.len() as u16).to_be_bytes());
    o.extend_from_slice(&AVC_PPS);
    o
}

/// AVC access unit `i`'s NAL units: an IDR slice when `key`, a non-IDR one otherwise.
pub fn avc_nals(i: u64, key: bool) -> Vec<Vec<u8>> {
    let mut n = vec![if key { 0x65 } else { 0x41 }];
    n.extend(body(i, if key { 2000 } else { 400 + (i % 5) as usize * 30 }));
    vec![n]
}

pub fn hevc_sets() -> Vec<Vec<u8>> {
    vec![vec![0x40, 0x01, 0x0c, 0x01], vec![0x42, 0x01, 0x01, 0x01, 0x60], vec![0x44, 0x01, 0xc1, 0x72]]
}

/// An HEVCDecoderConfigurationRecord: Main profile, level 3.1 (93), four-byte lengths, its
/// VPS, SPS and PPS.
pub fn hevc_record() -> Vec<u8> {
    let mut o = vec![1, 0x01, 0x60, 0, 0, 0, 0x90, 0, 0, 0, 0, 0, 93, 0xf0, 0, 0xfc, 0xfd, 0xf8, 0xf8, 0, 0, 0x0f, 3];
    for (ty, s) in [32u8, 33, 34].iter().zip(hevc_sets()) {
        o.push(0x80 | ty);
        o.extend_from_slice(&1u16.to_be_bytes());
        o.extend_from_slice(&(s.len() as u16).to_be_bytes());
        o.extend_from_slice(&s);
    }
    o
}

/// HEVC access unit `i`: an IDR_W_RADL slice (type 19) when `key`, TRAIL_R (1) otherwise.
pub fn hevc_nals(i: u64, key: bool) -> Vec<Vec<u8>> {
    let mut n = if key { vec![0x26, 0x01] } else { vec![0x02, 0x01] };
    n.extend(body(i + 50, if key { 1800 } else { 350 }));
    vec![n]
}

struct Bw(Vec<u8>, usize);
impl Bw {
    fn put(&mut self, v: u64, n: usize) {
        for i in (0..n).rev() {
            if self.1.is_multiple_of(8) {
                self.0.push(0);
            }
            if (v >> i) & 1 == 1 {
                *self.0.last_mut().unwrap() |= 0x80 >> (self.1 % 8);
            }
            self.1 += 1;
        }
    }
}

fn obu(ty: u8, payload: &[u8]) -> Vec<u8> {
    let mut o = vec![(ty << 3) | 0x02];
    let mut n = payload.len();
    loop {
        let b = (n & 0x7f) as u8;
        n >>= 7;
        if n == 0 {
            o.push(b);
            break;
        }
        o.push(b | 0x80);
    }
    o.extend_from_slice(payload);
    o
}

/// An AV1 sequence header OBU: profile 0, one operating point at level 8, the size.
pub fn av1_sequence_header(w: u32, h: u32) -> Vec<u8> {
    let mut b = Bw(Vec::new(), 0);
    b.put(0, 3);
    b.put(0, 1);
    b.put(0, 1);
    b.put(0, 1);
    b.put(0, 1);
    b.put(0, 5);
    b.put(0, 12);
    b.put(8, 5);
    b.put(0, 1);
    b.put(15, 4);
    b.put(15, 4);
    b.put(w as u64 - 1, 16);
    b.put(h as u64 - 1, 16);
    obu(1, &b.0)
}

/// An AV1CodecConfigurationRecord with the sequence header as its configOBUs.
pub fn av1_record(w: u32, h: u32) -> Vec<u8> {
    let mut o = vec![0x81, 0x08, 0x0c, 0x00];
    o.extend(av1_sequence_header(w, h));
    o
}

/// AV1 temporal unit `i`: a frame OBU (with a temporal delimiter before it, which the
/// gateway drops and its framing restores).
pub fn av1_tu(i: u64, key: bool) -> Vec<u8> {
    let mut o = vec![0x12, 0x00];
    o.extend(obu(6, &body(i + 90, if key { 1500 } else { 300 })));
    o
}

/// VP9 frame `i`: an uncompressed header (keyframe with its sync code, colour config and
/// size; or an interframe) and bytes.
pub fn vp9_frame(i: u64, key: bool, w: u32, h: u32) -> Vec<u8> {
    let mut b = Bw(Vec::new(), 0);
    b.put(2, 2);
    b.put(0, 2);
    b.put(0, 1);
    b.put(if key { 0 } else { 1 }, 1);
    b.put(1, 1);
    b.put(0, 1);
    if key {
        b.put(0x49_83_42, 24);
        b.put(1, 3);
        b.put(0, 1);
        b.put(w as u64 - 1, 16);
        b.put(h as u64 - 1, 16);
    }
    let mut f = b.0;
    f.extend(body(i + 300, if key { 1200 } else { 250 }));
    f
}

/// A VPCodecConfigurationRecord: profile 0, level 3.1, 8 bits.
pub fn vp9_record() -> Vec<u8> {
    vec![0, 31, 0x80, 1, 1, 1, 0, 0]
}

/// An AudioSpecificConfig: AAC-LC, 48 kHz, stereo.
pub const AAC_ASC: [u8; 2] = [0x11, 0x90];

pub fn aac_raw(i: u64) -> Vec<u8> {
    body(i + 1000, 120 + (i % 3) as usize)
}

/// The ADTS frame the gateway builds for a raw frame under `AAC_ASC`: profile 1, index 3,
/// two channels.
pub fn adts(raw: &[u8]) -> Vec<u8> {
    let len = raw.len() + 7;
    let mut o = vec![0xff, 0xf1, (1 << 6) | (3 << 2), (2 << 6) | ((len >> 11) & 3) as u8, ((len >> 3) & 0xff) as u8, (((len & 7) << 5) as u8) | 0x1f, 0xfc];
    o.extend_from_slice(raw);
    o
}

pub fn opus_packet(i: u64) -> Vec<u8> {
    let mut f = vec![0xfc];
    f.extend(body(i + 7_000, 80 + (i % 7) as usize));
    f
}

/// An Opus ID header (RFC 7845 §5.1): two channels, mapping family 0.
pub fn opus_head() -> Vec<u8> {
    let mut o = b"OpusHead".to_vec();
    o.extend_from_slice(&[1, 2, 0x38, 0x01, 0x80, 0xbb, 0, 0, 0, 0, 0]);
    o
}

/// `fLaC` and a STREAMINFO block (not marked last): 48 kHz, two channels, 16 bits.
pub fn flac_header() -> Vec<u8> {
    let mut o = b"fLaC".to_vec();
    o.extend_from_slice(&[0x00, 0, 0, 34]);
    let mut si = vec![0x10, 0x00, 0x10, 0x00, 0, 0, 0, 0, 0, 0];
    // sample rate 48 000 (20 bits), channels − 1 = 1 (3 bits), bits − 1 = 15 (5 bits),
    // total samples 0 (36 bits).
    let rate = 48_000u64;
    let v: u64 = rate << 44 | 1 << 41 | 15 << 36;
    si.extend_from_slice(&v.to_be_bytes());
    si.extend_from_slice(&[0; 16]);
    o.extend(si);
    o
}

pub fn flac_frame(i: u64) -> Vec<u8> {
    let mut f = vec![0xff, 0xf8, 0x69, 0x18, i as u8];
    f.extend(body(i + 5_000, 200));
    f
}