//! The aggregate share statement each aggregator signs (`WQS1`, telemetry-aggregation §6), the
//! aggregate receipt the collector signs (`WQA1`, §7), their verification (§8) and a room's
//! cell from its receipts (§9).

use ed25519_dalek::{Signer, SigningKey};
use sha2::{Digest, Sha256};

use super::dap::{report_checksum, xor32, CollectionJobReq, Interval, TaskConfig, R, W, BATCH_TIME_INTERVAL, ROLE_HELPER, ROLE_LEADER};
use super::measure::{self, Task, DURATIONS, OUTCOMES, STALLS};
use super::Reason;

pub const STATEMENT_MAGIC: &[u8; 4] = b"WQS1";
pub const RECEIPT_MAGIC: &[u8; 4] = b"WQA1";
pub const VERSION: u8 = 1;
/// The domain each signature is made under: the label, then 0x01.
pub const STATEMENT_DOMAIN: &[u8] = b"wai:telemetry-aggregate/share\x01";
pub const RECEIPT_DOMAIN: &[u8] = b"wai:telemetry-aggregate/receipt\x01";

pub fn sha256(parts: &[&[u8]]) -> [u8; 32] {
    let mut h = Sha256::new();
    for p in parts {
        h.update(p);
    }
    h.finalize().into()
}

fn sign(key: &SigningKey, domain: &[u8], body: &[u8]) -> [u8; 64] {
    let mut m = domain.to_vec();
    m.extend_from_slice(body);
    key.sign(&m).to_bytes()
}

fn verify_sig(key: &[u8; 32], domain: &[u8], body: &[u8], sig: &[u8; 64]) -> Result<(), Reason> {
    let mut m = domain.to_vec();
    m.extend_from_slice(body);
    crate::sig::verify_strict(key, &m, sig).map_err(|e| match e {
        crate::sig::SigError::BadKey => Reason::BadKey,
        crate::sig::SigError::BadSignature => Reason::BadSignature,
    })
}

/// An aggregate share statement (§6.1).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Statement {
    pub role: u8,
    pub task_id: [u8; 32],
    /// SHA-256 of the CollectionJobReq the share answers.
    pub collection: [u8; 32],
    /// The batch: the batch selector's interval.
    pub batch: Interval,
    pub report_count: u64,
    pub checksum: [u8; 32],
    pub agg_share: Vec<u8>,
    pub signer: [u8; 32],
    pub sig: [u8; 64],
}

impl Statement {
    fn body(&self) -> Vec<u8> {
        let mut w = W::default();
        w.raw(STATEMENT_MAGIC).u8(VERSION).u8(self.role).raw(&self.task_id).raw(&self.collection);
        w.raw(&self.batch.encode()).u64(self.report_count).raw(&self.checksum).opaque(4, &self.agg_share).raw(&self.signer);
        w.0
    }

    /// The statement, signed by `key`.
    #[allow(clippy::too_many_arguments)]
    pub fn signed(key: &SigningKey, role: u8, task_id: [u8; 32], collection: &CollectionJobReq, batch: Interval, report_count: u64, checksum: [u8; 32], agg_share: Vec<u8>) -> Statement {
        let mut s = Statement {
            role,
            task_id,
            collection: sha256(&[&collection.encode()]),
            batch,
            report_count,
            checksum,
            agg_share,
            signer: key.verifying_key().to_bytes(),
            sig: [0; 64],
        };
        s.sig = sign(key, STATEMENT_DOMAIN, &s.body());
        s
    }

    /// Sign the statement's fields again with `key`, which also becomes its signer.
    pub fn resigned(mut self, key: &SigningKey) -> Statement {
        self.signer = key.verifying_key().to_bytes();
        self.sig = sign(key, STATEMENT_DOMAIN, &self.body());
        self
    }

    pub fn encode(&self) -> Vec<u8> {
        let mut b = self.body();
        b.extend_from_slice(&self.sig);
        b
    }

    /// Read a statement (§6.2): exactly its bytes, magic `WQS1`, version 1, role 2 or 3.
    pub fn decode(b: &[u8]) -> Result<Statement, Reason> {
        let bad = |_| Reason::StatementInvalid;
        let mut r = R::new(b);
        if r.take(4).map_err(bad)? != STATEMENT_MAGIC || r.u8().map_err(bad)? != VERSION {
            return Err(Reason::StatementInvalid);
        }
        let role = r.u8().map_err(bad)?;
        if role != ROLE_LEADER && role != ROLE_HELPER {
            return Err(Reason::StatementInvalid);
        }
        let s = Statement {
            role,
            task_id: r.arr().map_err(bad)?,
            collection: r.arr().map_err(bad)?,
            batch: Interval { start: r.u64().map_err(bad)?, duration: r.u64().map_err(bad)? },
            report_count: r.u64().map_err(bad)?,
            checksum: r.arr().map_err(bad)?,
            agg_share: r.opaque(4, 0, 0xFFFF_FFFF).map_err(bad)?,
            signer: r.arr().map_err(bad)?,
            sig: r.arr().map_err(bad)?,
        };
        r.end().map_err(bad)?;
        Ok(s)
    }
}

/// An aggregate receipt (§7.1).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Receipt {
    pub task: Vec<u8>,
    pub collection: Vec<u8>,
    /// CollectionJobResp's interval: the smallest interval holding the batch's reports.
    pub interval: Interval,
    pub leader: Vec<u8>,
    /// Empty when the receipt holds one statement.
    pub helper: Vec<u8>,
    pub result: Vec<u128>,
    pub signer: [u8; 32],
    pub sig: [u8; 64],
}

impl Receipt {
    fn body(&self) -> Vec<u8> {
        let mut w = W::default();
        w.raw(RECEIPT_MAGIC).u8(VERSION).opaque(2, &self.task).opaque(4, &self.collection).raw(&self.interval.encode());
        w.opaque(4, &self.leader).opaque(4, &self.helper).u16(self.result.len() as u16);
        for x in &self.result {
            w.raw(&x.to_be_bytes());
        }
        w.raw(&self.signer);
        w.0
    }

    pub fn sign(mut self, key: &SigningKey) -> Receipt {
        self.signer = key.verifying_key().to_bytes();
        self.sig = sign(key, RECEIPT_DOMAIN, &self.body());
        self
    }

    pub fn encode(&self) -> Vec<u8> {
        let mut b = self.body();
        b.extend_from_slice(&self.sig);
        b
    }

    /// Read a receipt (§7.2): `truncated` when it ends early, `malformed` for another magic or
    /// bytes past its end, `unknown_version` for another version.
    pub fn decode(b: &[u8]) -> Result<Receipt, Reason> {
        let t = |_| Reason::Truncated;
        let mut r = R::new(b);
        if r.take(4).map_err(t)? != RECEIPT_MAGIC {
            return Err(Reason::Malformed);
        }
        if r.u8().map_err(t)? != VERSION {
            return Err(Reason::UnknownVersion);
        }
        let task = r.opaque(2, 0, 0xFFFF).map_err(t)?;
        let collection = r.opaque(4, 0, 0xFFFF_FFFF).map_err(t)?;
        let interval = Interval { start: r.u64().map_err(t)?, duration: r.u64().map_err(t)? };
        let leader = r.opaque(4, 0, 0xFFFF_FFFF).map_err(t)?;
        let helper = r.opaque(4, 0, 0xFFFF_FFFF).map_err(t)?;
        let n = r.u16().map_err(t)?;
        let mut result = Vec::with_capacity(n as usize);
        for _ in 0..n {
            result.push(u128::from_be_bytes(r.arr().map_err(t)?));
        }
        let rc = Receipt { task, collection, interval, leader, helper, result, signer: r.arr().map_err(t)?, sig: r.arr().map_err(t)? };
        r.end().map_err(|_| Reason::Malformed)?;
        Ok(rc)
    }

    /// The receipt's digest: SHA-256 over all its bytes but the signature.
    pub fn digest(&self) -> [u8; 32] {
        sha256(&[&self.body()])
    }
}

/// A report of the batch, as a reader is given it: its id, its time, and the task it was
/// uploaded to.
#[derive(Clone, Debug)]
pub struct GivenReport {
    pub id: [u8; 16],
    pub time: u64,
    pub task_id: [u8; 32],
}

/// What an accepted receipt establishes (§8.2).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Accepted {
    pub digest: [u8; 32],
    pub task_id: [u8; 32],
    pub measure: u8,
    pub vdaf: &'static str,
    pub vdaf_config: Vec<u8>,
    pub scope: [u8; 32],
    pub time_precision: u64,
    pub min_batch_size: u64,
    pub batch: Interval,
    pub interval: Interval,
    pub report_count: u64,
    pub result: Vec<u128>,
    pub leader: [u8; 32],
    pub helper: [u8; 32],
    pub collector: [u8; 32],
}

/// A refusal (§8.1): its code and where it is raised.
pub type Refusal = (&'static str, &'static str);

/// Verify a receipt (§8.1), with the batch's reports where they are given.
pub fn verify(bytes: &[u8], reports: Option<&[GivenReport]>) -> Result<Accepted, Refusal> {
    verify_with(bytes, reports, &|t| TaskConfig::decode(t).map_err(|_| Reason::TaskInvalid.code()).and_then(Task::new))
}

/// §8.1 steps 1 to 11 with step 2 given as `task`: from the receipt's `task` bytes, the task
/// that meets §4.3, or the code that refuses it. A task of measure 5 is checked against its
/// catalogue this way (impression-aggregates §8.1).
pub fn verify_with(bytes: &[u8], reports: Option<&[GivenReport]>, task: &dyn Fn(&[u8]) -> Result<Task, &'static str>) -> Result<Accepted, Refusal> {
    let at = |place: &'static str| move |r: Reason| (r.code(), place);
    // 1. The receipt reads.
    let rc = Receipt::decode(bytes).map_err(at("receipt"))?;
    // 2. The task meets §4.3.
    let task = task(&rc.task).map_err(|r| (r, "task"))?;
    // 3. The collector signed the receipt.
    if rc.signer != task.info.collector {
        return Err((Reason::WrongSigner.code(), "receipt"));
    }
    verify_sig(&rc.signer, RECEIPT_DOMAIN, &rc.body(), &rc.sig).map_err(at("receipt"))?;
    // 4. The collection job reads and is one this extension collects.
    let coll = CollectionJobReq::decode(&rc.collection).map_err(|_| (Reason::CollectionInvalid.code(), "collection"))?;
    let query = coll.interval().map_err(|_| (Reason::CollectionInvalid.code(), "collection"))?;
    if coll.batch_mode != BATCH_TIME_INTERVAL || query.duration == 0 || query.end().is_none() || !coll.agg_param.is_empty() || !coll.extensions.is_empty() {
        return Err((Reason::CollectionInvalid.code(), "collection"));
    }
    // 5. Two statements, by the task's two aggregators.
    if rc.helper.is_empty() || rc.leader.is_empty() {
        return Err((Reason::AggregatorAlone.code(), "statements"));
    }
    let l = Statement::decode(&rc.leader).map_err(at("leader"))?;
    let h = Statement::decode(&rc.helper).map_err(at("helper"))?;
    if l.signer == h.signer || l.role == h.role {
        return Err((Reason::AggregatorAlone.code(), "statements"));
    }
    for (s, want_role, want_key, at) in [(&l, ROLE_LEADER, task.info.leader, "leader"), (&h, ROLE_HELPER, task.info.helper, "helper")] {
        if s.role != want_role || s.signer != want_key {
            return Err((Reason::WrongSigner.code(), at));
        }
        verify_sig(&s.signer, STATEMENT_DOMAIN, &s.body(), &s.sig).map_err(|r| (r.code(), at))?;
    }
    // 6. Both statements are about this task and this collection.
    let coll_digest = sha256(&[&rc.collection]);
    for (s, at) in [(&l, "leader"), (&h, "helper")] {
        if s.task_id != task.id || s.collection != coll_digest {
            return Err((Reason::TaskMismatch.code(), at));
        }
    }
    // 7. The batch is within the query, and the reports' interval within the batch.
    for (s, at) in [(&l, "leader"), (&h, "helper")] {
        if s.batch.duration == 0 || !s.batch.within(&query) {
            return Err((Reason::BatchInvalid.code(), at));
        }
    }
    if !rc.interval.within(&l.batch) || rc.interval.duration == 0 {
        return Err((Reason::BatchInvalid.code(), "interval"));
    }
    // 8. The statements agree.
    if l.batch != h.batch || l.report_count != h.report_count || l.checksum != h.checksum {
        return Err((Reason::StatementsDisagree.code(), "statements"));
    }
    // 9. The batch is at least the task's minimum.
    if l.report_count < task.config.min_batch_size {
        return Err((Reason::BatchTooSmall.code(), "statements"));
    }
    // 10. The result is the shares', unsharded.
    for (s, at) in [(&l, "leader"), (&h, "helper")] {
        if !task.vdaf.is_agg_share(&s.agg_share) {
            return Err((Reason::ShareInvalid.code(), at));
        }
    }
    let result = task.vdaf.unshard(&[l.agg_share.clone(), h.agg_share.clone()]).map_err(|_| (Reason::ShareInvalid.code(), "helper"))?;
    if result != rc.result {
        return Err((Reason::ResultMismatch.code(), "result"));
    }
    // 11. The batch's reports, where given.
    if let Some(reps) = reports {
        let mut checksum = [0u8; 32];
        let mut ids = std::collections::BTreeSet::new();
        for g in reps {
            if g.task_id != task.id || !rc.interval.contains(g.time) || !ids.insert(g.id) {
                return Err((Reason::BatchMismatch.code(), "reports"));
            }
            xor32(&mut checksum, &report_checksum(&g.id));
        }
        if reps.len() as u64 != l.report_count || checksum != l.checksum {
            return Err((Reason::BatchMismatch.code(), "reports"));
        }
    }
    Ok(Accepted {
        digest: rc.digest(),
        task_id: task.id,
        measure: task.info.measure,
        vdaf: task.vdaf.name(),
        vdaf_config: task.config.vdaf_config.clone(),
        scope: task.info.scope,
        time_precision: task.config.time_precision,
        min_batch_size: task.config.min_batch_size,
        batch: l.batch,
        interval: rc.interval,
        report_count: l.report_count,
        result,
        leader: task.info.leader,
        helper: task.info.helper,
        collector: task.info.collector,
    })
}

/// A room's QoE cell (§9.2) from the accepted receipts of its `outcomes`, `stalls` and
/// `durations` tasks: live-metering §4.7's figures, each with the receipt it comes from.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Cell {
    pub scope: [u8; 32],
    pub batch: Interval,
    pub started: u128,
    pub start_failed: u128,
    pub exit_before_start: u128,
    pub stalls: u128,
    pub stall_ms: u128,
    pub engaged_ms: u128,
    pub vst_count: u128,
    pub vst_ms: u128,
    /// The report count of each receipt, in the order outcomes, stalls, durations.
    pub counts: [u64; 3],
}

pub fn cell(receipts: &[Accepted]) -> Result<Cell, Reason> {
    let get = |m: u8| -> Result<&Accepted, Reason> {
        let mut it = receipts.iter().filter(|a| a.measure == m);
        let a = it.next().ok_or(Reason::CellMismatch)?;
        if it.next().is_some() { Err(Reason::CellMismatch) } else { Ok(a) }
    };
    if receipts.len() != 3 {
        return Err(Reason::CellMismatch);
    }
    let (o, s, d) = (get(OUTCOMES)?, get(STALLS)?, get(DURATIONS)?);
    if o.scope == [0; 32] || [s, d].iter().any(|x| x.scope != o.scope || x.batch != o.batch) {
        return Err(Reason::CellMismatch);
    }
    Ok(Cell {
        scope: o.scope,
        batch: o.batch,
        started: o.result[0],
        start_failed: o.result[1],
        exit_before_start: o.result[2],
        vst_count: o.result[3],
        stalls: s.result[0],
        stall_ms: d.result[0],
        engaged_ms: d.result[1],
        vst_ms: d.result[2],
        counts: [o.report_count, s.report_count, d.report_count],
    })
}

/// The name of a measure, for reports.
pub fn measure_name(code: u8) -> &'static str {
    measure::name(code).unwrap_or("unknown")
}
