//! A reference client, leader, helper and collector (telemetry-aggregation §5, §6), each the
//! role `draft-ietf-ppm-dap-20` §4 gives it, in memory: upload (§4.4), one aggregation job
//! (§4.5.3), batch buckets (§4.5.4.3), collection (§4.6.1), the helper's aggregate share
//! (§4.6.4) and its encryption (§4.6.7); and, at each aggregate share, the statement the
//! aggregator signs and seals to the collector. Any party may run either aggregator; nothing
//! here holds both shares of a report.

use std::collections::{BTreeMap, BTreeSet};

use ed25519_dalek::SigningKey;

use super::dap::*;
use super::hpke;
use super::measure::Task;
use super::receipt::Statement;

/// An aggregator's keys: its HPKE configuration and private key, and its statement key.
pub struct Keys {
    pub hpke: HpkeConfig,
    pub hpke_sk: [u8; 32],
    pub signing: SigningKey,
}

/// A DAP problem type (§3.6) an aggregator aborts with.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Problem {
    InvalidMessage,
    BatchInvalid,
    InvalidBatchSize,
    InvalidAggregationParameter,
    BatchMismatch,
    BatchOverlap,
    UnsupportedExtension,
}

impl Problem {
    pub fn name(self) -> &'static str {
        match self {
            Problem::InvalidMessage => "invalidMessage",
            Problem::BatchInvalid => "batchInvalid",
            Problem::InvalidBatchSize => "invalidBatchSize",
            Problem::InvalidAggregationParameter => "invalidAggregationParameter",
            Problem::BatchMismatch => "batchMismatch",
            Problem::BatchOverlap => "batchOverlap",
            Problem::UnsupportedExtension => "unsupportedExtension",
        }
    }
}

/// The HPKE info string a statement is sealed to the collector under (§6.3): the label, the
/// sender's role and the collector's.
pub fn statement_info(role: u8) -> Vec<u8> {
    let mut v = b"wai-telemetry-aggregate statement".to_vec();
    v.extend_from_slice(&[role, ROLE_COLLECTOR]);
    v
}

/// What an aggregator accepts of a report's extensions (`draft-ietf-ppm-dap-20` §4.5.3.4), after
/// its input share decrypts: `Ok`, or the report error it rejects the report with and, where a
/// WAI text names one, the WAI reason it records beside it. A rule is `Sync` and
/// `RefUnwindSafe`, so an aggregator that holds one is `Send`, `Sync` and unwind-safe.
pub trait ReportRule: Sync + std::panic::RefUnwindSafe {
    fn extensions(&self, m: &ReportMetadata, private: &Extensions) -> Result<(), (u8, Option<&'static str>)>;
}

/// The rule of §5.1: a report carries no extension, public or private; any is
/// `invalid_message` (DAP §4.5.3.4 step 4).
pub struct NoExtensions;

impl ReportRule for NoExtensions {
    fn extensions(&self, m: &ReportMetadata, private: &Extensions) -> Result<(), (u8, Option<&'static str>)> {
        if !private.is_empty() || !m.public_extensions.is_empty() {
            return Err((report_error::INVALID_MESSAGE, None));
        }
        Ok(())
    }
}

/// The configuration a client seals an input share to, from its aggregator's HpkeConfigList:
/// the first, in the list's order of preference (DAP §4.4.1), of DAP §7's suite —
/// DHKEM(X25519, HKDF-SHA256), HKDF-SHA256 and AES-128-GCM — with a 32-byte key.
pub fn choose(list: &[HpkeConfig]) -> Option<&HpkeConfig> {
    list.iter().find(|c| c.kem == hpke::KEM_X25519_HKDF_SHA256 && c.kdf == hpke::KDF_HKDF_SHA256 && c.aead == hpke::AEAD_AES_128_GCM && c.public_key.len() == 32)
}

/// The client (§5.1): shard an encoded measurement and encrypt each input share to its
/// aggregator, under the configuration [`choose`] takes from its list. `rand` is Prio3's
/// randomness; `ikm` the seeds of the two ephemeral keys; `public` the report's public
/// extensions and `private` each input share's.
#[allow(clippy::too_many_arguments)]
pub fn client_report(
    task: &Task,
    leader: &[HpkeConfig],
    helper: &[HpkeConfig],
    id: [u8; 16],
    time: u64,
    encoded: &[u128],
    rand: &[u8],
    ikm: [&[u8]; 2],
    public: &Extensions,
    private: [&Extensions; 2],
) -> Report {
    let (public_share, shares) = task.vdaf.shard(&task.ctx(), encoded, &id, rand).expect("shard");
    let metadata = ReportMetadata { id, time, public_extensions: public.clone() };
    let aad = input_share_aad(&task.id, &task.config, &metadata, &public_share);
    let seal = |list: &[HpkeConfig], role: u8, share: &[u8], ikm: &[u8], ext: &Extensions| {
        let cfg = choose(list).expect("a configuration of DAP's suite");
        let pk: [u8; 32] = cfg.public_key.as_slice().try_into().expect("X25519 key");
        let (enc, payload) = hpke::seal_base(&pk, &input_share_info(role), &aad, &encode_plaintext_input_share(ext, share), ikm).expect("seal");
        HpkeCiphertext { config_id: cfg.id, enc, payload }
    };
    let leader_share = seal(leader, ROLE_LEADER, &shares[0], ikm[0], private[0]);
    let helper_share = seal(helper, ROLE_HELPER, &shares[1], ikm[1], private[1]);
    Report { metadata, public_share, leader_share, helper_share }
}

/// A batch bucket (§4.5.4.3), with the reports committed to it.
#[derive(Clone, Debug, Default)]
pub struct Bucket {
    pub agg: Vec<u8>,
    pub count: u64,
    pub checksum: [u8; 32],
    pub committed: Vec<([u8; 16], Vec<u8>)>,
}

/// Buckets merged: aggregate share, count, checksum, and the first and last bucket that hold
/// a report.
type Merged = (Vec<u8>, u64, [u8; 32], Option<(u64, u64)>);

/// One aggregator's report state: its task, keys, verification key, the rule it holds a
/// report's extensions to, buckets, the reports it has aggregated, the buckets it has
/// collected, and the WAI reason it recorded beside a report error.
pub struct Aggregator<'a> {
    pub task: &'a Task,
    pub keys: &'a Keys,
    pub verify_key: [u8; 32],
    pub rule: &'a dyn ReportRule,
    pub buckets: BTreeMap<u64, Bucket>,
    pub aggregated: BTreeSet<[u8; 16]>,
    pub collected: BTreeSet<u64>,
    pub reasons: BTreeMap<[u8; 16], &'static str>,
}

/// A report's outcome in an aggregation job: committed, or rejected with a report error.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Outcome {
    Committed,
    Rejected(u8),
}

impl<'a> Aggregator<'a> {
    pub fn new(task: &'a Task, keys: &'a Keys, verify_key: [u8; 32]) -> Self {
        Aggregator::with_rule(task, keys, verify_key, &NoExtensions)
    }

    /// An aggregator that holds a report's extensions to `rule`.
    pub fn with_rule(task: &'a Task, keys: &'a Keys, verify_key: [u8; 32], rule: &'a dyn ReportRule) -> Self {
        Aggregator {
            task,
            keys,
            verify_key,
            rule,
            buckets: BTreeMap::new(),
            aggregated: BTreeSet::new(),
            collected: BTreeSet::new(),
            reasons: BTreeMap::new(),
        }
    }

    /// A report error, with the WAI reason recorded beside it where there is one.
    fn reject(&mut self, id: &[u8; 16], e: (u8, Option<&'static str>)) -> u8 {
        if let Some(r) = e.1 {
            self.reasons.insert(*id, r);
        }
        e.0
    }

    /// Whether an output share may be committed (§4.5.4.3).
    fn eligible(&self, m: &ReportMetadata) -> Result<(), u8> {
        if self.collected.contains(&m.time) {
            return Err(report_error::BATCH_COLLECTED);
        }
        if self.aggregated.contains(&m.id) {
            return Err(report_error::REPORT_REPLAYED);
        }
        Ok(())
    }

    fn commit(&mut self, m: &ReportMetadata, out: Vec<u8>) {
        let b = self.buckets.entry(m.time).or_insert_with(|| Bucket { agg: self.task.vdaf.agg_init(), ..Default::default() });
        b.agg = self.task.vdaf.merge(&b.agg, &out).expect("output share");
        b.count += 1;
        xor32(&mut b.checksum, &report_checksum(&m.id));
        b.committed.push((m.id, out));
        self.aggregated.insert(m.id);
    }

    /// Decrypt and check this aggregator's input share of a report (§4.5.3.3, §4.5.3.4): its
    /// extensions are held to the aggregator's rule.
    fn open(&self, role: u8, m: &ReportMetadata, public_share: &[u8], ct: &HpkeCiphertext) -> Result<Vec<u8>, (u8, Option<&'static str>)> {
        if ct.config_id != self.keys.hpke.id {
            return Err((report_error::HPKE_UNKNOWN_CONFIG_ID, None));
        }
        let aad = input_share_aad(&self.task.id, &self.task.config, m, public_share);
        let pt = hpke::open_base(&ct.enc, &self.keys.hpke_sk, &input_share_info(role), &aad, &ct.payload).map_err(|_| (report_error::HPKE_DECRYPT_ERROR, None))?;
        let (exts, payload) = decode_plaintext_input_share(&pt).map_err(|_| (report_error::INVALID_MESSAGE, None))?;
        self.rule.extensions(m, &exts)?;
        if self.task.config.task_interval().is_some_and(|i| !i.contains(m.time)) {
            return Err((report_error::REPORT_DROPPED, None));
        }
        Ok(payload)
    }

    /// The buckets of a batch interval, merged: aggregate share, count, checksum, and the
    /// first and last bucket holding a report.
    fn merged(&self, i: &Interval) -> Merged {
        let mut agg = self.task.vdaf.agg_init();
        let (mut count, mut checksum, mut span) = (0, [0u8; 32], None::<(u64, u64)>);
        for (t, b) in self.buckets.range(i.start..i.end().unwrap_or(u64::MAX)) {
            agg = self.task.vdaf.merge(&agg, &b.agg).expect("bucket");
            count += b.count;
            xor32(&mut checksum, &b.checksum);
            span = Some(span.map_or((*t, *t), |(a, _)| (a, *t)));
        }
        (agg, count, checksum, span)
    }

    fn overlaps_collected(&self, i: &Interval) -> bool {
        self.collected.range(i.start..i.end().unwrap_or(u64::MAX)).next().is_some()
    }

    /// Seal an aggregate share and a statement of it to the collector (§4.6.7, §6.3).
    fn seal_share(&self, role: u8, collector: &HpkeConfig, coll: &CollectionJobReq, st: &Statement, ikm: [&[u8]; 2]) -> (HpkeCiphertext, HpkeCiphertext) {
        let pk: [u8; 32] = collector.public_key.as_slice().try_into().expect("X25519 key");
        let aad = aggregate_share_aad(&self.task.id, &self.task.config, coll);
        let (enc, payload) = hpke::seal_base(&pk, &aggregate_share_info(role), &aad, &st.agg_share, ikm[0]).expect("seal");
        let share = HpkeCiphertext { config_id: collector.id, enc, payload };
        let (enc, payload) = hpke::seal_base(&pk, &statement_info(role), &aad, &st.encode(), ikm[1]).expect("seal");
        (share, HpkeCiphertext { config_id: collector.id, enc, payload })
    }
}

/// The leader (§5.2): its uploads and the reports it has not yet aggregated.
pub struct Leader<'a> {
    pub agg: Aggregator<'a>,
    pub pending: Vec<Report>,
    pub seen: BTreeSet<[u8; 16]>,
}

/// One aggregation job, as the leader ran it.
pub struct Job {
    pub init_req: Vec<u8>,
    pub resp: Vec<u8>,
    /// Each report the leader started with, and its outcome.
    pub outcomes: Vec<([u8; 16], Outcome)>,
}

impl<'a> Leader<'a> {
    pub fn new(agg: Aggregator<'a>) -> Self {
        Leader { agg, pending: Vec::new(), seen: BTreeSet::new() }
    }

    /// Handle an upload request (§4.4.2.2): the reports refused, each with its error.
    pub fn upload(&mut self, body: &[u8]) -> Result<Vec<([u8; 16], u8)>, Problem> {
        let reports = decode_upload(body).map_err(|_| Problem::InvalidMessage)?;
        let mut errors = Vec::new();
        for r in reports {
            let m = &r.metadata;
            if r.leader_share.config_id != self.agg.keys.hpke.id {
                errors.push((m.id, report_error::HPKE_UNKNOWN_CONFIG_ID));
            } else if !self.seen.insert(m.id) || self.agg.collected.contains(&m.time) {
                errors.push((m.id, report_error::REPORT_REPLAYED));
            } else if self.agg.task.config.task_interval().is_some_and(|i| !i.contains(m.time)) {
                errors.push((m.id, report_error::REPORT_DROPPED));
            } else {
                self.pending.push(r);
            }
        }
        Ok(errors)
    }

    /// Run one aggregation job over every pending report with the helper (§4.5.3).
    pub fn aggregate(&mut self, helper: &mut Helper) -> Result<Job, Problem> {
        let task = self.agg.task;
        let ctx = task.ctx();
        let pending = std::mem::take(&mut self.pending);
        let mut outcomes = Vec::new();
        let mut inits = Vec::new();
        let mut states = Vec::new();
        for r in &pending {
            let m = &r.metadata;
            let opened = self.agg.eligible(m).map_err(|e| (e, None)).and_then(|_| self.agg.open(ROLE_LEADER, m, &r.public_share, &r.leader_share));
            let opened = opened.map_err(|e| self.agg.reject(&m.id, e));
            let started = opened.and_then(|payload| {
                task.vdaf.leader_init(&self.agg.verify_key, &ctx, &m.id, &r.public_share, &payload).map_err(|_| report_error::VDAF_VERIFY_ERROR)
            });
            match started {
                Ok((state, outbound)) => {
                    inits.push(VerifyInit { metadata: m.clone(), public_share: r.public_share.clone(), encrypted: r.helper_share.clone(), payload: outbound });
                    states.push((m.clone(), state));
                }
                Err(e) => outcomes.push((m.id, Outcome::Rejected(e))),
            }
        }
        let req = AggregationJobInitReq { verification_key_id: 0, agg_param: vec![], extensions: vec![], inits };
        let init_req = req.encode();
        let resp = helper.aggregation_job_init(&init_req)?;
        let resps = decode_job_resp(&resp).map_err(|_| Problem::InvalidMessage)?;
        if resps.len() != states.len() || resps.iter().zip(&states).any(|(v, (m, _))| v.report_id != m.id) {
            return Err(Problem::InvalidMessage);
        }
        for (v, (m, state)) in resps.into_iter().zip(states) {
            let o = match v.kind {
                VerifyRespKind::Continue(payload) => match task.vdaf.leader_continued(&state, &payload) {
                    Ok(out) => {
                        self.agg.commit(&m, out);
                        Outcome::Committed
                    }
                    Err(_) => Outcome::Rejected(report_error::VDAF_VERIFY_ERROR),
                },
                VerifyRespKind::Reject(e) => Outcome::Rejected(e),
                VerifyRespKind::Finish => return Err(Problem::InvalidMessage),
            };
            outcomes.push((m.id, o));
        }
        Ok(Job { init_req, resp, outcomes })
    }

    /// Handle a collection job (§4.6.1) with the helper (§4.6.4): the response, the helper's
    /// sealed statement and the leader's own, and the AggregateShareReq it sent. At `now_s`, in
    /// seconds, a job whose batch interval has not ended `hold_s` seconds before is still
    /// processing (`Ok(None)`): a task that holds no collection passes a hold of 0.
    #[allow(clippy::too_many_arguments)]
    pub fn collect(&mut self, body: &[u8], helper: &mut Helper, collector: &HpkeConfig, ikm: [&[u8]; 4], now_s: u64, hold_s: u64) -> Result<Option<Collection>, Problem> {
        let task = self.agg.task;
        let coll = CollectionJobReq::decode(body).map_err(|_| Problem::InvalidMessage)?;
        if coll.batch_mode != task.config.batch_mode {
            return Err(Problem::InvalidMessage);
        }
        if !coll.agg_param.is_empty() {
            return Err(Problem::InvalidAggregationParameter);
        }
        let query = coll.interval().map_err(|_| Problem::InvalidMessage)?;
        if query.duration == 0 || query.end().is_none() {
            return Err(Problem::BatchInvalid);
        }
        if self.agg.overlaps_collected(&query) {
            return Err(Problem::BatchOverlap);
        }
        if !coll.extensions.is_empty() {
            return Err(Problem::UnsupportedExtension);
        }
        let due = query.end().and_then(|e| e.checked_mul(task.config.time_precision)).and_then(|e| e.checked_add(hold_s));
        if due.is_none_or(|d| now_s < d) {
            return Ok(None);
        }
        let (agg, count, checksum, span) = self.agg.merged(&query);
        if count < task.config.min_batch_size {
            return Err(Problem::InvalidBatchSize);
        }
        let req = AggregateShareReq { collection: coll.clone(), selector_mode: BATCH_TIME_INTERVAL, selector_config: query.encode(), report_count: count, checksum };
        let share_req = req.encode();
        let (helper_share, helper_statement) = helper.aggregate_share(&share_req, collector, [ikm[2], ikm[3]])?;
        let st = Statement::signed(&self.agg.keys.signing, ROLE_LEADER, task.id, &coll, query, count, checksum, agg);
        let (leader_share, leader_statement) = self.agg.seal_share(ROLE_LEADER, collector, &coll, &st, [ikm[0], ikm[1]]);
        for t in query.start..query.end().unwrap() {
            self.agg.collected.insert(t);
        }
        let (a, b) = span.expect("a batch of at least the minimum holds a report");
        let resp = CollectionJobResp { report_count: count, interval: Interval { start: a, duration: b - a + 1 }, leader_share, helper_share: HpkeCiphertext::decode(&helper_share).map_err(|_| Problem::InvalidMessage)? };
        Ok(Some(Collection { request: body.to_vec(), share_req, helper_share, resp: resp.encode(), leader_statement: leader_statement.encode(), helper_statement }))
    }
}

/// A collection, as the leader ran it.
pub struct Collection {
    pub request: Vec<u8>,
    pub share_req: Vec<u8>,
    /// The helper's AggregateShare: an HpkeCiphertext.
    pub helper_share: Vec<u8>,
    pub resp: Vec<u8>,
    /// Each aggregator's statement, sealed to the collector (HpkeCiphertext).
    pub leader_statement: Vec<u8>,
    pub helper_statement: Vec<u8>,
}

/// The helper (§5.2).
pub struct Helper<'a> {
    pub agg: Aggregator<'a>,
}

impl<'a> Helper<'a> {
    /// Handle an AggregationJobInitReq (§4.5.3.2): the AggregationJobResp.
    pub fn aggregation_job_init(&mut self, body: &[u8]) -> Result<Vec<u8>, Problem> {
        let task = self.agg.task;
        let ctx = task.ctx();
        let req = AggregationJobInitReq::decode(body).map_err(|_| Problem::InvalidMessage)?;
        if !req.extensions.is_empty() {
            return Err(Problem::UnsupportedExtension);
        }
        if !req.agg_param.is_empty() {
            return Err(Problem::InvalidAggregationParameter);
        }
        let ids: BTreeSet<[u8; 16]> = req.inits.iter().map(|v| v.metadata.id).collect();
        if ids.len() != req.inits.len() {
            return Err(Problem::InvalidMessage);
        }
        let mut resps = Vec::new();
        for v in &req.inits {
            let m = &v.metadata;
            let r = if req.verification_key_id != 0 {
                Err(report_error::UNKNOWN_VERIFICATION_KEY_ID)
            } else {
                let opened = self.agg.open(ROLE_HELPER, m, &v.public_share, &v.encrypted).and_then(|p| self.agg.eligible(m).map(|_| p).map_err(|e| (e, None)));
                opened
                    .map_err(|e| self.agg.reject(&m.id, e))
                    .and_then(|p| task.vdaf.helper_init(&self.agg.verify_key, &ctx, &m.id, &v.public_share, &p, &v.payload).map_err(|_| report_error::VDAF_VERIFY_ERROR))
            };
            let kind = match r {
                Ok((out, outbound)) => {
                    self.agg.commit(m, out);
                    VerifyRespKind::Continue(outbound)
                }
                Err(e) => VerifyRespKind::Reject(e),
            };
            resps.push(VerifyResp { report_id: m.id, kind });
        }
        Ok(encode_job_resp(&resps))
    }

    /// Handle an AggregateShareReq (§4.6.4): the AggregateShare (an HpkeCiphertext) and the
    /// helper's statement sealed to the collector.
    pub fn aggregate_share(&mut self, body: &[u8], collector: &HpkeConfig, ikm: [&[u8]; 2]) -> Result<(Vec<u8>, Vec<u8>), Problem> {
        let task = self.agg.task;
        let req = AggregateShareReq::decode(body).map_err(|_| Problem::InvalidMessage)?;
        if req.selector_mode != task.config.batch_mode || req.collection.batch_mode != task.config.batch_mode {
            return Err(Problem::InvalidMessage);
        }
        let query = req.collection.interval().map_err(|_| Problem::InvalidMessage)?;
        let batch = Interval::decode(&req.selector_config).map_err(|_| Problem::InvalidMessage)?;
        if batch.duration == 0 || !batch.within(&query) {
            return Err(Problem::BatchInvalid);
        }
        if self.agg.overlaps_collected(&batch) {
            return Err(Problem::BatchOverlap);
        }
        let (agg, count, checksum, _) = self.agg.merged(&batch);
        if count < task.config.min_batch_size {
            return Err(Problem::InvalidBatchSize);
        }
        if !req.collection.agg_param.is_empty() {
            return Err(Problem::InvalidMessage);
        }
        if !req.collection.extensions.is_empty() {
            return Err(Problem::UnsupportedExtension);
        }
        if count != req.report_count || checksum != req.checksum {
            return Err(Problem::BatchMismatch);
        }
        let st = Statement::signed(&self.agg.keys.signing, ROLE_HELPER, task.id, &req.collection, batch, count, checksum, agg);
        let (share, statement) = self.agg.seal_share(ROLE_HELPER, collector, &req.collection, &st, ikm);
        for t in batch.start..batch.end().unwrap() {
            self.agg.collected.insert(t);
        }
        Ok((share.encode(), statement.encode()))
    }
}

/// The collector (§5.3): open both aggregate shares and both statements, require each
/// statement's share to be the share DAP delivered, and unshard.
pub struct Opened {
    pub result: Vec<u128>,
    pub leader: Statement,
    pub helper: Statement,
}

pub fn collector_open(task: &Task, sk: &[u8; 32], coll: &CollectionJobReq, resp: &[u8], statements: [&[u8]; 2]) -> Result<Opened, &'static str> {
    let resp = CollectionJobResp::decode(resp).map_err(|_| "collection job response")?;
    let aad = aggregate_share_aad(&task.id, &task.config, coll);
    let open = |ct: &HpkeCiphertext, info: Vec<u8>| hpke::open_base(&ct.enc, sk, &info, &aad, &ct.payload).map_err(|_| "decryption");
    let l_share = open(&resp.leader_share, aggregate_share_info(ROLE_LEADER))?;
    let h_share = open(&resp.helper_share, aggregate_share_info(ROLE_HELPER))?;
    let mut sts = Vec::new();
    for (b, role) in [(statements[0], ROLE_LEADER), (statements[1], ROLE_HELPER)] {
        let ct = HpkeCiphertext::decode(b).map_err(|_| "sealed statement")?;
        sts.push(Statement::decode(&open(&ct, statement_info(role))?).map_err(|_| "statement")?);
    }
    let h = sts.pop().unwrap();
    let l = sts.pop().unwrap();
    if l.agg_share != l_share || h.agg_share != h_share || l.report_count != resp.report_count || h.report_count != resp.report_count {
        return Err("a statement is not of the share delivered");
    }
    let result = task.vdaf.unshard(&[l_share, h_share]).map_err(|_| "unshard")?;
    Ok(Opened { result, leader: l, helper: h })
}

#[cfg(test)]
mod tests {
    use super::*;

    /// An aggregator, and so a leader and a helper, is `Send`, `Sync` and unwind-safe whatever
    /// report rule it holds: the trait's bounds carry those of the rule.
    #[test]
    fn an_aggregator_is_send_sync_and_unwind_safe() {
        fn holds<T: Send + Sync + std::panic::UnwindSafe + std::panic::RefUnwindSafe>() {}
        holds::<Aggregator<'static>>();
        holds::<Leader<'static>>();
        holds::<Helper<'static>>();
    }
}
