//! The QoE measures (extensions/telemetry-aggregation.md §3) and the task a measure runs under
//! (§4): what each measure reads from one session's metrics (delivery-telemetry §7.2), the
//! Prio3 variant and configuration that aggregates it, and the WAI layout of DAP's
//! `task_info`.

use sha2::{Digest, Sha256};

use super::dap::{TaskConfig, BATCH_TIME_INTERVAL, EXT_TASK_INTERVAL};
use super::prio3::Vdaf;
use super::Reason;
use crate::telemetry::qoe::SessionMetrics;

/// The batch floor (§4.3): a task's minimum batch size is at least this, as delivery-telemetry
/// §8.4's floor is 50 sessions.
pub const FLOOR: u64 = 50;

pub const OUTCOMES: u8 = 1;
pub const STALLS: u8 = 2;
pub const DURATIONS: u8 = 3;
pub const START_TIME: u8 = 4;
/// Ad presentations counted over a bearer with no return path (impression-aggregates §5.1):
/// its VDAF is not fixed here but by the catalogue whose digest is the task's `scope`.
pub const IMPRESSIONS: u8 = 5;

/// impression-aggregates §5.3's code, which a task of measure 5 is refused with when its
/// `measure` or `scope` is not its catalogue's, and when no catalogue is given (§4.3 rule 3).
pub const CATALOGUE_MISMATCH: &str = "catalogue_mismatch";

/// The largest stall count a session reports to `stalls`.
pub const STALLS_MAX: u64 = 65_535;
/// The largest duration, in milliseconds, a session reports to `durations`.
pub const DURATION_MAX: u64 = u32::MAX as u64;
/// The upper edges, in milliseconds, of the first six `start_time` buckets; the seventh holds
/// every later start and the eighth a session that never presented.
pub const START_EDGES_MS: [u64; 6] = [250, 500, 1000, 2000, 4000, 8000];

/// The measures of §3: code, name.
pub const MEASURES: [(u8, &str); 5] =
    [(OUTCOMES, "outcomes"), (STALLS, "stalls"), (DURATIONS, "durations"), (START_TIME, "start_time"), (IMPRESSIONS, "impressions")];

pub fn name(code: u8) -> Option<&'static str> {
    MEASURES.iter().find(|(c, _)| *c == code).map(|(_, n)| *n)
}

/// The VDAF a measure is aggregated with (§3), two aggregators; `None` for measure 5, whose
/// VDAF its catalogue fixes, and for a code of no measure.
pub fn vdaf(code: u8) -> Option<Vdaf> {
    Some(match code {
        OUTCOMES => Vdaf::sumvec(2, 4, 1, 2),
        STALLS => Vdaf::sum(2, STALLS_MAX),
        DURATIONS => Vdaf::sumvec(2, 3, DURATION_MAX, 10),
        START_TIME => Vdaf::histogram(2, 8, 3),
        _ => return None,
    })
}

/// A session's measurement for a measure (§3), in the form `Vdaf::encode` takes.
pub fn measurement(code: u8, m: &SessionMetrics) -> Option<Vec<u128>> {
    let b = |x: bool| x as u128;
    Some(match code {
        OUTCOMES => vec![b(m.started), b(m.start_failed), b(m.exit_before_start), b(m.vst_ms.is_some())],
        STALLS => vec![m.stalls.min(STALLS_MAX) as u128],
        DURATIONS => [m.stall_ms, m.engaged_ms, m.vst_ms.unwrap_or(0)].iter().map(|&x| x.min(DURATION_MAX) as u128).collect(),
        START_TIME => vec![match m.vst_ms {
            None => 7,
            Some(v) => START_EDGES_MS.iter().position(|&e| v < e).unwrap_or(6) as u128,
        }],
        _ => return None,
    })
}

/// The label that opens a WAI task's `task_info` (§4.1).
pub const LABEL: &[u8; 8] = b"wai-qoe1";
pub const TASK_INFO_LEN: usize = 8 + 1 + 32 + 16 + 32 * 3;

/// DAP's `task_info` as this extension lays it out (§4.1).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TaskInfo {
    pub measure: u8,
    /// Zero for a service's own task; the room id for a room's (§9).
    pub scope: [u8; 32],
    pub salt: [u8; 16],
    pub leader: [u8; 32],
    pub helper: [u8; 32],
    pub collector: [u8; 32],
}

impl TaskInfo {
    pub fn encode(&self) -> Vec<u8> {
        let mut v = LABEL.to_vec();
        v.push(self.measure);
        v.extend_from_slice(&self.scope);
        v.extend_from_slice(&self.salt);
        v.extend_from_slice(&self.leader);
        v.extend_from_slice(&self.helper);
        v.extend_from_slice(&self.collector);
        v
    }

    pub fn decode(b: &[u8]) -> Option<TaskInfo> {
        if b.len() != TASK_INFO_LEN || &b[..8] != LABEL {
            return None;
        }
        let a = |i: usize| -> [u8; 32] { b[i..i + 32].try_into().unwrap() };
        Some(TaskInfo { measure: b[8], scope: a(9), salt: b[41..57].try_into().unwrap(), leader: a(57), helper: a(89), collector: a(121) })
    }
}

/// The task id (§4.2): SHA-256 of `wai:telemetry-aggregate/task`, a zero byte and the
/// TaskConfiguration.
pub fn task_id(cfg: &TaskConfig) -> [u8; 32] {
    let mut h = Sha256::new();
    h.update(b"wai:telemetry-aggregate/task\x00");
    h.update(cfg.encode());
    h.finalize().into()
}

/// A task that meets §4.3: its configuration, its id, its `task_info` and its VDAF.
pub struct Task {
    pub config: TaskConfig,
    pub id: [u8; 32],
    pub info: TaskInfo,
    pub vdaf: Vdaf,
}

/// A rule 3 for a measure whose VDAF §3.1 does not fix (measure 5): given the configuration
/// and its `task_info`, the VDAF, or the code that refuses it.
pub type MeasureRule<'a> = &'a dyn Fn(&TaskConfig, &TaskInfo) -> Result<Vdaf, &'static str>;

/// §4.3 rule 1: the configuration's `task_info`, or `task_invalid`.
pub fn rule_1(config: &TaskConfig) -> Result<TaskInfo, &'static str> {
    let invalid = Reason::TaskInvalid.code();
    let info = TaskInfo::decode(&config.task_info).ok_or(invalid)?;
    if config.time_precision == 0 || config.batch_mode != BATCH_TIME_INTERVAL || !config.batch_config.is_empty() {
        return Err(invalid);
    }
    let ext_ok = config.extensions.iter().all(|(t, d)| *t == EXT_TASK_INTERVAL && super::dap::Interval::decode(d).is_ok());
    if !ext_ok || !super::dap::exts_sorted(&config.extensions) {
        return Err(invalid);
    }
    Ok(info)
}

impl Task {
    /// Check a configuration against §4.3, in order, with `rule` as rule 3 for measure 5
    /// (impression-aggregates §5.3 steps 2 to 4). A refusal is its code.
    pub fn with(config: TaskConfig, rule: MeasureRule<'_>) -> Result<Task, &'static str> {
        let info = rule_1(&config)?;
        if name(info.measure).is_none() {
            return Err(Reason::MeasureUnknown.code());
        }
        let vdaf = match vdaf(info.measure) {
            Some(v) => {
                if config.vdaf_type != v.id() || config.vdaf_config != v.dap_config() {
                    return Err(Reason::MeasureMismatch.code());
                }
                v
            }
            None => rule(&config, &info)?,
        };
        if config.min_batch_size < FLOOR {
            return Err(Reason::FloorTooLow.code());
        }
        if info.leader == info.helper {
            return Err(Reason::AggregatorAlone.code());
        }
        Ok(Task { id: task_id(&config), config, info, vdaf })
    }

    /// Check a configuration against §4.3 with no catalogue: measures 1 to 4 as §3.1 fixes
    /// them, and measure 5 refused `catalogue_mismatch`.
    pub fn new(config: TaskConfig) -> Result<Task, &'static str> {
        Task::with(config, &|_, _| Err(CATALOGUE_MISMATCH))
    }

    /// The VDAF application context DAP shards and verifies under: `dap-20` and the task id.
    pub fn ctx(&self) -> Vec<u8> {
        let mut c = super::dap::DAP_TAG.to_vec();
        c.extend_from_slice(&self.id);
        c
    }
}
