//! A `stage` claim from what a sink observed (`extensions/staged-delivery.md` §10.2,
//! jwp-receipts §2.4.10).
//!
//! A [`StageObserver`] watches one stage set at a sink: the layer envelopes that arrived,
//! with the evidence of their delivery receipts; the layer objects that never arrived and
//! are known only by their receipts (dropped at a relay, abandoned at a deadline,
//! jwp-receipts §4.4); and every [`OfferReport`] the sink gave for them. With the base's
//! chain it states the claim's body: each object's status, its first reason and its
//! evidence; the units by stage, the conflicts and the deadline misses; the tier and
//! stage hash of what is presented; and the stage the delivered objects reach. The body
//! is signed with [`ClaimReceipt::seal`](crate::claim_receipt::ClaimReceipt::seal), whose
//! checks are §2.4.10's rules, so a body this module states and a verifier refuses never
//! leaves the sink signed.

use super::{OfferReport, ReceiptEvidence, StageChain, StageRefusal};
use crate::claim_receipt::{stage_point, StageClaim, StageEvidence, StageObject, StageStatus, STAGE_VOCABULARY};
use crate::codecs::DeterminismTier;
use crate::container::{RefForm, Refines, Wai};
use std::collections::{BTreeMap, BTreeSet};

/// Why an envelope could not be observed as a layer of the set.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ObserveError {
    /// Not a `WAI1` envelope, or one with no `refines` that reads.
    NotALayer,
    /// A layer of another base.
    OtherBase,
    /// An event of an epoch earlier than one the observer has already seen.
    StaleEpoch,
}

/// What became of one layer object, as the sink saw it.
#[derive(Debug, Clone)]
struct Seen {
    layer: u8,
    /// The units it covers; `None` for every unit of the base (no `span`).
    span: Option<(u32, u32)>,
    /// Its envelope bytes; 0 for an object that never arrived.
    bytes: u64,
    arrived: bool,
    evidence: ReceiptEvidence,
    envelope: Option<StageRefusal>,
    held_for_base: bool,
    applied: BTreeSet<u32>,
    refused: BTreeMap<u32, StageRefusal>,
    held: BTreeSet<u32>,
}

impl Seen {
    fn report(&mut self, r: &OfferReport) {
        // A report decides for the envelope as of now: a refusal replaces what was held,
        // and a report that refuses nothing for the envelope, or applies a unit, clears an
        // earlier refusal (a layer refused, then offered again and applied).
        match r.envelope {
            Some(why) => (self.envelope, self.held_for_base) = (Some(why), false),
            None => (self.envelope, self.held_for_base) = (None, r.held_for_base),
        }
        if !r.applied.is_empty() {
            (self.envelope, self.held_for_base) = (None, false);
        }
        for u in &r.applied {
            self.applied.insert(*u);
            self.refused.remove(u);
            self.held.remove(u);
        }
        for (u, why) in &r.refused {
            if !self.applied.contains(u) {
                self.refused.insert(*u, *why);
                self.held.remove(u);
            }
        }
        for u in &r.held {
            if !self.applied.contains(u) && !self.refused.contains_key(u) {
                self.held.insert(*u);
            }
        }
    }

    /// The status, and its reason, over the `count` units the object covers. A unit
    /// still held when the claim is made counts as not applied, for the reason it is
    /// held for (`layer-gap`, or `base-absent` for the whole envelope).
    fn status(&self, from: u32, count: u32) -> (StageStatus, Option<String>) {
        if !self.arrived {
            return (StageStatus::Dropped, None);
        }
        // What the sink applied decides first: a unit applied stays applied, whatever a
        // later offer of the same envelope was told.
        if self.applied.is_empty() {
            if let Some(r) = self.envelope {
                return (StageStatus::Refused, Some(r.code().into()));
            }
            if self.held_for_base {
                return (StageStatus::Refused, Some(StageRefusal::BaseAbsent.code().into()));
            }
        }
        let units = from..from.saturating_add(count);
        let not_applied: Vec<(u32, StageRefusal)> = units
            .filter(|u| !self.applied.contains(u))
            .map(|u| (u, self.refused.get(&u).copied().unwrap_or(StageRefusal::LayerGap)))
            .collect();
        match (not_applied.first(), not_applied.len() as u64 == u64::from(count)) {
            (None, _) => (StageStatus::Applied, None),
            (Some(_), true) if not_applied.iter().all(|(_, r)| *r == StageRefusal::DeadlineMissed) => (StageStatus::Late, None),
            (Some((_, r)), true) => (StageStatus::Refused, Some(r.code().into())),
            (Some((_, r)), false) => (StageStatus::Partial, Some(r.code().into())),
        }
    }
}

fn evidence_of(e: ReceiptEvidence) -> StageEvidence {
    match e {
        ReceiptEvidence::DeliveredVerified => StageEvidence::DeliveredVerified,
        ReceiptEvidence::InclusionVerified => StageEvidence::InclusionVerified,
        ReceiptEvidence::None => StageEvidence::None,
    }
}

/// One stage set at a sink, observed for a `stage` claim.
#[derive(Debug, Clone)]
pub struct StageObserver {
    base: [u8; 32],
    base_bytes: u64,
    base_delivered: bool,
    requested: Option<u8>,
    committed: Option<u8>,
    objects: BTreeMap<[u8; 32], Seen>,
    /// The sink's epoch for the base ([`super::StageSink::base_epoch`]) at the last event;
    /// `None` before any.
    epoch: Option<u64>,
}

impl StageObserver {
    /// A set whose base has content hash `base` and `base_bytes` envelope bytes.
    /// `base_delivered` is whether the sink verified the base's delivery receipt over its
    /// bytes before it offered the base's first layer: only then is a stage stated as
    /// delivered (jwp-receipts §2.4.10, `op_delivered`).
    pub fn new(base: [u8; 32], base_bytes: u64, base_delivered: bool) -> StageObserver {
        StageObserver { base, base_bytes, base_delivered, requested: None, committed: None, objects: BTreeMap::new(), epoch: None }
    }

    /// Whether the sink verified the base's delivery receipt over its bytes before it
    /// offered the base's first layer in the current epoch. A new epoch clears it.
    pub fn set_base_delivered(&mut self, delivered: bool) {
        self.base_delivered = delivered;
    }

    /// Every event carries the sink's epoch for the base ([`super::StageSink::base_epoch`]).
    /// A later epoch means the sink forgot the base (staged-delivery §4.4, §14): decoded
    /// again, it is a base the sink never decoded and never attested, and nothing it did
    /// before is part of what it presents. The observation starts over: an envelope the
    /// sink still holds for the base keeps its arrival, its bytes and its evidence, and is
    /// held for the base again with no decision; every other object is dropped; and no
    /// delivered stage is stated until the caller says otherwise. An event of an earlier
    /// epoch than one already seen is stale, and is ignored (`false`).
    fn at(&mut self, epoch: u64) -> bool {
        match self.epoch {
            Some(e) if epoch < e => return false,
            Some(e) if epoch > e => self.restart(),
            _ => {}
        }
        self.epoch = Some(epoch);
        true
    }

    fn restart(&mut self) {
        self.objects.retain(|_, s| s.arrived && s.held_for_base);
        for s in self.objects.values_mut() {
            (s.envelope, s.held_for_base) = (None, true);
            s.applied.clear();
            s.refused.clear();
            s.held.clear();
        }
        self.base_delivered = false;
    }

    /// The stage the sink asked for and the stage the serving party said it would
    /// deliver. They are stated only beside a delivered stage.
    pub fn request(&mut self, requested: Option<u8>, committed: Option<u8>) {
        (self.requested, self.committed) = (requested, committed);
    }

    /// A layer envelope that arrived, with the evidence of its delivery receipt, before
    /// the sink offers it. Arriving again keeps the stronger evidence (staged-delivery §4.3).
    pub fn arrived(&mut self, epoch: u64, envelope: &[u8], evidence: ReceiptEvidence) -> Result<(), ObserveError> {
        let w = Wai::from_bytes(envelope).map_err(|_| ObserveError::NotALayer)?;
        let r = w.manifest.refines.as_ref().ok_or(ObserveError::NotALayer)?;
        let r = Refines::from_value(r, RefForm::ContentHash).map_err(|_| ObserveError::NotALayer)?;
        if r.base != hex(&self.base) {
            return Err(ObserveError::OtherBase);
        }
        if !self.at(epoch) {
            return Err(ObserveError::StaleEpoch);
        }
        let hash = crate::jwp_hash::content_hash(envelope);
        let seen = self.objects.entry(hash).or_insert_with(|| Seen {
            layer: r.layer,
            span: r.span.map(|s| (s.from, s.count)),
            bytes: 0,
            arrived: false,
            evidence,
            envelope: None,
            held_for_base: false,
            applied: BTreeSet::new(),
            refused: BTreeMap::new(),
            held: BTreeSet::new(),
        });
        // The envelope's own `refines`, which its content hash authenticates, says what
        // it is, whatever an earlier receipt-only sighting supposed (`dropped`).
        seen.layer = r.layer;
        seen.span = r.span.map(|s| (s.from, s.count));
        seen.arrived = true;
        seen.bytes = envelope.len() as u64;
        if rank(evidence) > rank(seen.evidence) {
            seen.evidence = evidence;
        }
        Ok(())
    }

    /// A layer object that never arrived, known by its object receipt: its content hash,
    /// its layer, the units it covers (`None` for every unit) and what was verified of
    /// the receipt — by inclusion, or nothing (jwp-receipts §4.4). An object that arrived
    /// is not made dropped by this.
    pub fn dropped(&mut self, epoch: u64, content_hash: [u8; 32], layer: u8, span: Option<(u32, u32)>, evidence: ReceiptEvidence) {
        if !self.at(epoch) {
            return;
        }
        let evidence = if evidence == ReceiptEvidence::DeliveredVerified { ReceiptEvidence::InclusionVerified } else { evidence };
        self.objects.entry(content_hash).or_insert_with(|| Seen {
            layer,
            span,
            bytes: 0,
            arrived: false,
            evidence,
            envelope: None,
            held_for_base: false,
            applied: BTreeSet::new(),
            refused: BTreeMap::new(),
            held: BTreeSet::new(),
        });
    }

    /// The reports the sink gave — on an offer, a release or an expiry — with the sink's
    /// epoch for the base after it gave them. A report about an object this observer has
    /// not seen arrive is not this set's, and is ignored.
    pub fn reports(&mut self, epoch: u64, reports: &[OfferReport]) {
        if !self.at(epoch) {
            return;
        }
        for r in reports {
            if let Some(s) = self.objects.get_mut(&r.content_hash).filter(|s| s.arrived) {
                s.report(r);
            }
        }
    }

    /// The claim's body as of now, over `chain`, the base's chain at the sink, whose epoch
    /// for the base is `epoch`. A claim never spans two epochs: if the sink has forgotten
    /// the base since the last event, the claim is made over an observation started over.
    pub fn claim(&self, chain: &StageChain, epoch: u64) -> StageClaim {
        if self.epoch.is_some_and(|e| epoch > e) {
            let mut fresh = self.clone();
            fresh.restart();
            fresh.epoch = Some(epoch);
            return fresh.claim(chain, epoch);
        }
        let n = chain.n_units();
        let mut objects: Vec<StageObject> = self
            .objects
            .iter()
            .map(|(hash, s)| {
                let (from, count) = s.span.unwrap_or((0, n));
                let (status, reason) = s.status(from, count);
                StageObject {
                    layer: u64::from(s.layer),
                    object: *hash,
                    units: (u64::from(from), u64::from(count)),
                    bytes: s.bytes,
                    status,
                    evidence: evidence_of(s.evidence),
                    reason,
                }
            })
            .collect();
        objects.sort_by_key(|o| (o.layer, o.units.0, o.object));
        let mut units_by_stage = vec![0u64; 1 + chain.units().iter().map(|u| usize::from(u.stage)).max().unwrap_or(0)];
        for u in chain.units() {
            units_by_stage[usize::from(u.stage)] += 1;
        }
        let lowest = units_by_stage.iter().position(|c| *c != 0).unwrap_or(0) as u64;
        // A unit presented below the stage its layers would have reached on time: one a
        // layer missed the deadline of (§8).
        let missed: BTreeSet<u32> = self
            .objects
            .values()
            .flat_map(|s| s.refused.iter().filter(|(_, r)| **r == StageRefusal::DeadlineMissed).map(|(u, _)| *u))
            .collect();
        let tiers: Vec<Option<DeterminismTier>> = (0..n).map(|u| chain.tier(u)).collect();
        let tier = if tiers.iter().all(|t| *t == Some(DeterminismTier::DecodeEquivalence)) {
            "decode-equivalence"
        } else if tiers.iter().all(|t| t.is_some()) {
            "entropy-consistency"
        } else {
            "none"
        };
        let mut c = StageClaim {
            vocabulary: STAGE_VOCABULARY.into(),
            op_requested: None,
            op_committed: None,
            op_delivered: None,
            presented: stage_point(lowest),
            units_by_stage,
            tier: tier.into(),
            stage_hash: chain.clip_sha256(),
            bytes_envelopes: self.base_bytes + objects.iter().map(|o| o.bytes).sum::<u64>(),
            objects,
            conflicts: chain.units().iter().filter(|u| u.frozen).count() as u64,
            deadline_misses: missed.len() as u64,
        };
        if self.base_delivered {
            c.op_delivered = Some(stage_point(c.delivered_stage()));
            c.op_requested = self.requested.map(|l| stage_point(u64::from(l)));
            c.op_committed = self.committed.map(|l| stage_point(u64::from(l)));
        }
        c
    }
}

/// Evidence ordered by strength (staged-delivery §4.3).
fn rank(e: ReceiptEvidence) -> u8 {
    match e {
        ReceiptEvidence::DeliveredVerified => 2,
        ReceiptEvidence::InclusionVerified => 1,
        ReceiptEvidence::None => 0,
    }
}

fn hex(h: &[u8; 32]) -> String {
    h.iter().map(|b| format!("{b:02x}")).collect()
}
