//! The messages of `draft-ietf-ppm-dap-20` (§4) this extension uses, in the TLS presentation
//! language, each written and read strictly: the task configuration (§4.2), HPKE configurations
//! and ciphertexts (§4.1, §4.4.1), reports and their input shares (§4.4.2), aggregation jobs
//! (§4.5.3), collection jobs and aggregate shares (§4.6), in the time-interval batch mode (§5.1).

/// A DAP-20 domain separation prefix (`"dap-20"`); the drafts replace it with `"dap"` on
/// publication.
pub const DAP_TAG: &[u8] = b"dap-20";

pub const ROLE_COLLECTOR: u8 = 0;
pub const ROLE_CLIENT: u8 = 1;
pub const ROLE_LEADER: u8 = 2;
pub const ROLE_HELPER: u8 = 3;

pub const BATCH_TIME_INTERVAL: u8 = 1;
pub const BATCH_LEADER_SELECTED: u8 = 2;
pub const EXT_TASK_INTERVAL: u16 = 1;

/// The report errors of §4.1.
pub mod report_error {
    pub const BATCH_COLLECTED: u8 = 1;
    pub const REPORT_REPLAYED: u8 = 2;
    pub const REPORT_DROPPED: u8 = 3;
    pub const HPKE_UNKNOWN_CONFIG_ID: u8 = 4;
    pub const HPKE_DECRYPT_ERROR: u8 = 5;
    pub const VDAF_VERIFY_ERROR: u8 = 6;
    pub const INVALID_MESSAGE: u8 = 7;
    pub const REPORT_TOO_EARLY: u8 = 8;
    pub const UNKNOWN_VERIFICATION_KEY_ID: u8 = 9;
    pub const UNSUPPORTED_EXTENSION: u8 = 10;

    /// The name §4.1 gives a report error.
    pub fn name(e: u8) -> &'static str {
        match e {
            BATCH_COLLECTED => "batch_collected",
            REPORT_REPLAYED => "report_replayed",
            REPORT_DROPPED => "report_dropped",
            HPKE_UNKNOWN_CONFIG_ID => "hpke_unknown_config_id",
            HPKE_DECRYPT_ERROR => "hpke_decrypt_error",
            VDAF_VERIFY_ERROR => "vdaf_verify_error",
            INVALID_MESSAGE => "invalid_message",
            REPORT_TOO_EARLY => "report_too_early",
            UNKNOWN_VERIFICATION_KEY_ID => "unknown_verification_key_id",
            UNSUPPORTED_EXTENSION => "unsupported_extension",
            _ => "reserved",
        }
    }
}

/// A message that does not read.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Malformed;

/// Writes the presentation language.
#[derive(Default)]
pub struct W(pub Vec<u8>);

impl W {
    pub fn u8(&mut self, v: u8) -> &mut Self {
        self.0.push(v);
        self
    }
    pub fn u16(&mut self, v: u16) -> &mut Self {
        self.0.extend_from_slice(&v.to_be_bytes());
        self
    }
    pub fn u32(&mut self, v: u32) -> &mut Self {
        self.0.extend_from_slice(&v.to_be_bytes());
        self
    }
    pub fn u64(&mut self, v: u64) -> &mut Self {
        self.0.extend_from_slice(&v.to_be_bytes());
        self
    }
    pub fn raw(&mut self, v: &[u8]) -> &mut Self {
        self.0.extend_from_slice(v);
        self
    }
    /// An opaque vector with a length prefix of `n` bytes.
    pub fn opaque(&mut self, n: usize, v: &[u8]) -> &mut Self {
        let len = v.len() as u64;
        self.0.extend_from_slice(&len.to_be_bytes()[8 - n..]);
        self.0.extend_from_slice(v);
        self
    }
}

/// Reads the presentation language; every read past the end is `Malformed`.
pub struct R<'a> {
    b: &'a [u8],
}

impl<'a> R<'a> {
    pub fn new(b: &'a [u8]) -> R<'a> {
        R { b }
    }
    pub fn take(&mut self, n: usize) -> Result<&'a [u8], Malformed> {
        if self.b.len() < n {
            return Err(Malformed);
        }
        let (h, t) = self.b.split_at(n);
        self.b = t;
        Ok(h)
    }
    pub fn u8(&mut self) -> Result<u8, Malformed> {
        Ok(self.take(1)?[0])
    }
    pub fn u16(&mut self) -> Result<u16, Malformed> {
        Ok(u16::from_be_bytes(self.take(2)?.try_into().unwrap()))
    }
    pub fn u32(&mut self) -> Result<u32, Malformed> {
        Ok(u32::from_be_bytes(self.take(4)?.try_into().unwrap()))
    }
    pub fn u64(&mut self) -> Result<u64, Malformed> {
        Ok(u64::from_be_bytes(self.take(8)?.try_into().unwrap()))
    }
    pub fn arr<const N: usize>(&mut self) -> Result<[u8; N], Malformed> {
        Ok(self.take(N)?.try_into().unwrap())
    }
    /// An opaque vector with a length prefix of `n` bytes and a length in `min..=max`.
    pub fn opaque(&mut self, n: usize, min: usize, max: u64) -> Result<Vec<u8>, Malformed> {
        let mut w = [0u8; 8];
        w[8 - n..].copy_from_slice(self.take(n)?);
        let len = u64::from_be_bytes(w);
        if len > max || (len as usize) < min {
            return Err(Malformed);
        }
        Ok(self.take(len as usize)?.to_vec())
    }
    pub fn is_empty(&self) -> bool {
        self.b.is_empty()
    }
    pub fn end(&self) -> Result<(), Malformed> {
        if self.b.is_empty() { Ok(()) } else { Err(Malformed) }
    }
}

/// An extension list (§4.2.2, §4.4.3, §4.5.1, §4.6.2): `(type, data)` pairs.
pub type Extensions = Vec<(u16, Vec<u8>)>;

fn write_exts(w: &mut W, e: &Extensions) {
    let mut inner = W::default();
    for (t, d) in e {
        inner.u16(*t).opaque(2, d);
    }
    w.opaque(2, &inner.0);
}

fn read_exts(r: &mut R) -> Result<Extensions, Malformed> {
    let body = r.opaque(2, 0, 0xFFFF)?;
    let mut ir = R::new(&body);
    let mut out = Vec::new();
    while !ir.is_empty() {
        out.push((ir.u16()?, ir.opaque(2, 0, 0xFFFF)?));
    }
    Ok(out)
}

/// Whether the extension types rise strictly (§4.2.2).
pub fn exts_sorted(e: &Extensions) -> bool {
    e.windows(2).all(|w| w[0].0 < w[1].0)
}

/// TaskConfiguration (§4.2).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TaskConfig {
    pub task_info: Vec<u8>,
    pub leader: Vec<u8>,
    pub helper: Vec<u8>,
    pub time_precision: u64,
    pub min_batch_size: u64,
    pub batch_mode: u8,
    pub batch_config: Vec<u8>,
    pub vdaf_type: u32,
    pub vdaf_config: Vec<u8>,
    pub extensions: Extensions,
}

impl TaskConfig {
    pub fn write(&self, w: &mut W) {
        w.opaque(1, &self.task_info).opaque(2, &self.leader).opaque(2, &self.helper).u64(self.time_precision).u64(self.min_batch_size);
        w.u8(self.batch_mode).opaque(2, &self.batch_config).u32(self.vdaf_type).opaque(2, &self.vdaf_config);
        write_exts(w, &self.extensions);
    }
    pub fn encode(&self) -> Vec<u8> {
        let mut w = W::default();
        self.write(&mut w);
        w.0
    }
    pub fn read(r: &mut R) -> Result<TaskConfig, Malformed> {
        Ok(TaskConfig {
            task_info: r.opaque(1, 0, 0xFF)?,
            leader: r.opaque(2, 1, 0xFFFF)?,
            helper: r.opaque(2, 1, 0xFFFF)?,
            time_precision: r.u64()?,
            min_batch_size: r.u64()?,
            batch_mode: r.u8()?,
            batch_config: r.opaque(2, 0, 0xFFFF)?,
            vdaf_type: r.u32()?,
            vdaf_config: r.opaque(2, 0, 0xFFFF)?,
            extensions: read_exts(r)?,
        })
    }
    pub fn decode(b: &[u8]) -> Result<TaskConfig, Malformed> {
        let mut r = R::new(b);
        let t = TaskConfig::read(&mut r)?;
        r.end()?;
        Ok(t)
    }
    /// The task_interval extension's interval, where the task has one.
    pub fn task_interval(&self) -> Option<Interval> {
        self.extensions.iter().find(|(t, _)| *t == EXT_TASK_INTERVAL).and_then(|(_, d)| Interval::decode(d).ok())
    }
}

/// HpkeConfig (§4.4.1).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct HpkeConfig {
    pub id: u8,
    pub kem: u16,
    pub kdf: u16,
    pub aead: u16,
    pub public_key: Vec<u8>,
}

impl HpkeConfig {
    pub fn encode(&self) -> Vec<u8> {
        let mut w = W::default();
        w.u8(self.id).u16(self.kem).u16(self.kdf).u16(self.aead).opaque(2, &self.public_key);
        w.0
    }
    pub fn decode(b: &[u8]) -> Result<HpkeConfig, Malformed> {
        let mut r = R::new(b);
        let c = HpkeConfig { id: r.u8()?, kem: r.u16()?, kdf: r.u16()?, aead: r.u16()?, public_key: r.opaque(2, 1, 0xFFFF)? };
        r.end()?;
        Ok(c)
    }
}

/// HpkeCiphertext (§4.1).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct HpkeCiphertext {
    pub config_id: u8,
    pub enc: Vec<u8>,
    pub payload: Vec<u8>,
}

impl HpkeCiphertext {
    pub fn write(&self, w: &mut W) {
        w.u8(self.config_id).opaque(2, &self.enc).opaque(4, &self.payload);
    }
    pub fn read(r: &mut R) -> Result<HpkeCiphertext, Malformed> {
        Ok(HpkeCiphertext { config_id: r.u8()?, enc: r.opaque(2, 1, 0xFFFF)?, payload: r.opaque(4, 1, 0xFFFF_FFFF)? })
    }
    pub fn encode(&self) -> Vec<u8> {
        let mut w = W::default();
        self.write(&mut w);
        w.0
    }
    pub fn decode(b: &[u8]) -> Result<HpkeCiphertext, Malformed> {
        let mut r = R::new(b);
        let c = HpkeCiphertext::read(&mut r)?;
        r.end()?;
        Ok(c)
    }
}

/// Interval (§4.1.1), in units of the task's time precision.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Interval {
    pub start: u64,
    pub duration: u64,
}

impl Interval {
    pub fn encode(&self) -> Vec<u8> {
        let mut w = W::default();
        w.u64(self.start).u64(self.duration);
        w.0
    }
    pub fn decode(b: &[u8]) -> Result<Interval, Malformed> {
        let mut r = R::new(b);
        let i = Interval { start: r.u64()?, duration: r.u64()? };
        r.end()?;
        Ok(i)
    }
    pub fn end(&self) -> Option<u64> {
        self.start.checked_add(self.duration)
    }
    pub fn contains(&self, t: u64) -> bool {
        self.end().is_some_and(|e| self.start <= t && t < e)
    }
    /// Whether `self` falls within `other` (§4.1.1).
    pub fn within(&self, other: &Interval) -> bool {
        match (self.end(), other.end()) {
            (Some(a), Some(b)) => self.start >= other.start && a <= b,
            _ => false,
        }
    }
}

/// ReportMetadata (§4.4.2).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ReportMetadata {
    pub id: [u8; 16],
    pub time: u64,
    pub public_extensions: Extensions,
}

impl ReportMetadata {
    pub fn write(&self, w: &mut W) {
        w.raw(&self.id).u64(self.time);
        write_exts(w, &self.public_extensions);
    }
    pub fn read(r: &mut R) -> Result<ReportMetadata, Malformed> {
        Ok(ReportMetadata { id: r.arr()?, time: r.u64()?, public_extensions: read_exts(r)? })
    }
}

/// Report (§4.4.2).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Report {
    pub metadata: ReportMetadata,
    pub public_share: Vec<u8>,
    pub leader_share: HpkeCiphertext,
    pub helper_share: HpkeCiphertext,
}

impl Report {
    pub fn write(&self, w: &mut W) {
        self.metadata.write(w);
        w.opaque(4, &self.public_share);
        self.leader_share.write(w);
        self.helper_share.write(w);
    }
    pub fn read(r: &mut R) -> Result<Report, Malformed> {
        Ok(Report {
            metadata: ReportMetadata::read(r)?,
            public_share: r.opaque(4, 0, 0xFFFF_FFFF)?,
            leader_share: HpkeCiphertext::read(r)?,
            helper_share: HpkeCiphertext::read(r)?,
        })
    }
    pub fn encode(&self) -> Vec<u8> {
        let mut w = W::default();
        self.write(&mut w);
        w.0
    }
    pub fn decode(b: &[u8]) -> Result<Report, Malformed> {
        let mut r = R::new(b);
        let x = Report::read(&mut r)?;
        r.end()?;
        Ok(x)
    }
}

/// UploadRequest (§4.4.2): reports, one after another, to the end of the message.
pub fn decode_upload(b: &[u8]) -> Result<Vec<Report>, Malformed> {
    let mut r = R::new(b);
    let mut out = Vec::new();
    while !r.is_empty() {
        out.push(Report::read(&mut r)?);
    }
    Ok(out)
}

pub fn encode_upload(reports: &[Report]) -> Vec<u8> {
    let mut w = W::default();
    for x in reports {
        x.write(&mut w);
    }
    w.0
}

/// PlaintextInputShare (§4.4.2).
pub fn encode_plaintext_input_share(private_extensions: &Extensions, payload: &[u8]) -> Vec<u8> {
    let mut w = W::default();
    write_exts(&mut w, private_extensions);
    w.opaque(4, payload);
    w.0
}

pub fn decode_plaintext_input_share(b: &[u8]) -> Result<(Extensions, Vec<u8>), Malformed> {
    let mut r = R::new(b);
    let e = read_exts(&mut r)?;
    let p = r.opaque(4, 1, 0xFFFF_FFFF)?;
    r.end()?;
    Ok((e, p))
}

/// InputShareAad (§4.4.2).
pub fn input_share_aad(task_id: &[u8; 32], cfg: &TaskConfig, m: &ReportMetadata, public_share: &[u8]) -> Vec<u8> {
    let mut w = W::default();
    w.raw(task_id);
    cfg.write(&mut w);
    m.write(&mut w);
    w.opaque(4, public_share);
    w.0
}

/// The HPKE info string of an input share (§4.4.2): the client's role and the recipient's.
pub fn input_share_info(server_role: u8) -> Vec<u8> {
    let mut v = DAP_TAG.to_vec();
    v.extend_from_slice(b" input share");
    v.extend_from_slice(&[ROLE_CLIENT, server_role]);
    v
}

/// The HPKE info string of an aggregate share (§4.6.7).
pub fn aggregate_share_info(server_role: u8) -> Vec<u8> {
    let mut v = DAP_TAG.to_vec();
    v.extend_from_slice(b" aggregate share");
    v.extend_from_slice(&[server_role, ROLE_COLLECTOR]);
    v
}

/// VerifyInit (§4.5.3.1): a report share and the leader's outbound message.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct VerifyInit {
    pub metadata: ReportMetadata,
    pub public_share: Vec<u8>,
    pub encrypted: HpkeCiphertext,
    pub payload: Vec<u8>,
}

/// AggregationJobInitReq (§4.5.3.1).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AggregationJobInitReq {
    pub verification_key_id: u8,
    pub agg_param: Vec<u8>,
    pub extensions: Extensions,
    pub inits: Vec<VerifyInit>,
}

impl AggregationJobInitReq {
    pub fn encode(&self) -> Vec<u8> {
        let mut w = W::default();
        w.u8(self.verification_key_id).opaque(4, &self.agg_param);
        write_exts(&mut w, &self.extensions);
        for v in &self.inits {
            v.metadata.write(&mut w);
            w.opaque(4, &v.public_share);
            v.encrypted.write(&mut w);
            w.opaque(4, &v.payload);
        }
        w.0
    }
    pub fn decode(b: &[u8]) -> Result<AggregationJobInitReq, Malformed> {
        let mut r = R::new(b);
        let verification_key_id = r.u8()?;
        let agg_param = r.opaque(4, 0, 0xFFFF_FFFF)?;
        let extensions = read_exts(&mut r)?;
        let mut inits = Vec::new();
        while !r.is_empty() {
            inits.push(VerifyInit {
                metadata: ReportMetadata::read(&mut r)?,
                public_share: r.opaque(4, 0, 0xFFFF_FFFF)?,
                encrypted: HpkeCiphertext::read(&mut r)?,
                payload: r.opaque(4, 1, 0xFFFF_FFFF)?,
            });
        }
        Ok(AggregationJobInitReq { verification_key_id, agg_param, extensions, inits })
    }
}

/// VerifyResp (§4.5.3.2).
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum VerifyRespKind {
    Continue(Vec<u8>),
    Finish,
    Reject(u8),
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct VerifyResp {
    pub report_id: [u8; 16],
    pub kind: VerifyRespKind,
}

pub fn encode_job_resp(resps: &[VerifyResp]) -> Vec<u8> {
    let mut w = W::default();
    for v in resps {
        w.raw(&v.report_id);
        match &v.kind {
            VerifyRespKind::Continue(p) => {
                w.u8(0).opaque(4, p);
            }
            VerifyRespKind::Finish => {
                w.u8(1);
            }
            VerifyRespKind::Reject(e) => {
                w.u8(2).u8(*e);
            }
        }
    }
    w.0
}

pub fn decode_job_resp(b: &[u8]) -> Result<Vec<VerifyResp>, Malformed> {
    let mut r = R::new(b);
    let mut out = Vec::new();
    while !r.is_empty() {
        let report_id = r.arr()?;
        let kind = match r.u8()? {
            0 => VerifyRespKind::Continue(r.opaque(4, 1, 0xFFFF_FFFF)?),
            1 => VerifyRespKind::Finish,
            2 => VerifyRespKind::Reject(r.u8()?),
            _ => return Err(Malformed),
        };
        out.push(VerifyResp { report_id, kind });
    }
    Ok(out)
}

/// CollectionJobReq (§4.6.1) with a Query of the time-interval mode (§5.1.1) or another.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CollectionJobReq {
    pub batch_mode: u8,
    pub query_config: Vec<u8>,
    pub agg_param: Vec<u8>,
    pub extensions: Extensions,
}

impl CollectionJobReq {
    pub fn time_interval(i: Interval) -> CollectionJobReq {
        CollectionJobReq { batch_mode: BATCH_TIME_INTERVAL, query_config: i.encode(), agg_param: vec![], extensions: vec![] }
    }
    pub fn write(&self, w: &mut W) {
        w.u8(self.batch_mode).opaque(2, &self.query_config).opaque(4, &self.agg_param);
        write_exts(w, &self.extensions);
    }
    pub fn read(r: &mut R) -> Result<CollectionJobReq, Malformed> {
        Ok(CollectionJobReq { batch_mode: r.u8()?, query_config: r.opaque(2, 0, 0xFFFF)?, agg_param: r.opaque(4, 0, 0xFFFF_FFFF)?, extensions: read_exts(r)? })
    }
    pub fn encode(&self) -> Vec<u8> {
        let mut w = W::default();
        self.write(&mut w);
        w.0
    }
    pub fn decode(b: &[u8]) -> Result<CollectionJobReq, Malformed> {
        let mut r = R::new(b);
        let c = CollectionJobReq::read(&mut r)?;
        r.end()?;
        Ok(c)
    }
    /// The batch interval of a time-interval query.
    pub fn interval(&self) -> Result<Interval, Malformed> {
        if self.batch_mode != BATCH_TIME_INTERVAL {
            return Err(Malformed);
        }
        Interval::decode(&self.query_config)
    }
}

/// AggregateShareReq (§4.6.4): the collection job, the batch selector, the report count and
/// the checksum.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AggregateShareReq {
    pub collection: CollectionJobReq,
    pub selector_mode: u8,
    pub selector_config: Vec<u8>,
    pub report_count: u64,
    pub checksum: [u8; 32],
}

impl AggregateShareReq {
    pub fn encode(&self) -> Vec<u8> {
        let mut w = W::default();
        self.collection.write(&mut w);
        w.u8(self.selector_mode).opaque(2, &self.selector_config).u64(self.report_count).raw(&self.checksum);
        w.0
    }
    pub fn decode(b: &[u8]) -> Result<AggregateShareReq, Malformed> {
        let mut r = R::new(b);
        let x = AggregateShareReq {
            collection: CollectionJobReq::read(&mut r)?,
            selector_mode: r.u8()?,
            selector_config: r.opaque(2, 0, 0xFFFF)?,
            report_count: r.u64()?,
            checksum: r.arr()?,
        };
        r.end()?;
        Ok(x)
    }
}

/// AggregateShareAad (§4.6.7).
pub fn aggregate_share_aad(task_id: &[u8; 32], cfg: &TaskConfig, c: &CollectionJobReq) -> Vec<u8> {
    let mut w = W::default();
    w.raw(task_id);
    cfg.write(&mut w);
    c.write(&mut w);
    w.0
}

/// CollectionJobResp (§4.6.1).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CollectionJobResp {
    pub report_count: u64,
    pub interval: Interval,
    pub leader_share: HpkeCiphertext,
    pub helper_share: HpkeCiphertext,
}

impl CollectionJobResp {
    pub fn encode(&self) -> Vec<u8> {
        let mut w = W::default();
        w.u64(self.report_count).raw(&self.interval.encode());
        self.leader_share.write(&mut w);
        self.helper_share.write(&mut w);
        w.0
    }
    pub fn decode(b: &[u8]) -> Result<CollectionJobResp, Malformed> {
        let mut r = R::new(b);
        let x = CollectionJobResp {
            report_count: r.u64()?,
            interval: Interval { start: r.u64()?, duration: r.u64()? },
            leader_share: HpkeCiphertext::read(&mut r)?,
            helper_share: HpkeCiphertext::read(&mut r)?,
        };
        r.end()?;
        Ok(x)
    }
}

/// The checksum a report adds to its batch bucket (§4.5.4.3): SHA-256 of its ID.
pub fn report_checksum(id: &[u8; 16]) -> [u8; 32] {
    use sha2::{Digest, Sha256};
    Sha256::digest(id).into()
}

pub fn xor32(a: &mut [u8; 32], b: &[u8; 32]) {
    for (x, y) in a.iter_mut().zip(b) {
        *x ^= y;
    }
}
