//! Authorization tokens: the alias cache a receiver keeps for each session
//! (`draft-ietf-moq-transport-22` §8.9, §9.1.3, §9.1.4), and the relay's grants.
//!
//! A grant names a namespace prefix, the rights it gives (publish, subscribe), and the
//! SHA-256 of the token value that carries it. Tokens are of type 0, agreed out of band
//! (§8.9: "Type 0 is reserved to indicate that the type is not defined in the table and
//! is negotiated out-of-band"). A grant rule with `proof` names a room grant's record,
//! which every participant of the room holds: it gives its rights only on a session whose
//! SETUP proved possession of the grant's principal's key (WAI live-rooms §10.5).
//!
//! A relay that serves a room whose authority issues grants as visitors arrive cannot name
//! each in advance: it is given the authority's key instead, and honours every grant that
//! key signed, mapped onto rights as live-rooms §10.2 maps it, beside its proof and within
//! its window. A field `*` in a rule's prefix matches any one field.

use crate::config::AuthConfig;
use joule_moq::draft22::{SessionError, Token};
use joule_moq::grant_proof::GrantHead;
use joule_moq::TrackNamespace;
use std::collections::{HashMap, HashSet};

/// A token's size against MAX_AUTH_TOKEN_CACHE_SIZE: 16 bytes and its value (§9.1.3).
fn size_of(value: &[u8]) -> u64 {
    16 + value.len() as u64
}

/// One token as a message resolves it: its type and value.
pub type Resolved = (u64, Vec<u8>);

/// The aliases one peer has registered with this endpoint.
#[derive(Debug)]
pub struct TokenCache {
    max: u64,
    used: u64,
    entries: HashMap<u64, Resolved>,
}

impl TokenCache {
    /// A cache bounded by the MAX_AUTH_TOKEN_CACHE_SIZE this endpoint declared.
    pub fn new(max: u64) -> Self {
        TokenCache { max, used: 0, entries: HashMap::new() }
    }

    pub fn used(&self) -> u64 {
        self.used
    }

    /// The tokens of the peer's SETUP. A REGISTER past the limit is read as USE_VALUE,
    /// never a failure (§9.1.4); a duplicate alias still closes the session. A client's
    /// SETUP carries no DELETE or USE_ALIAS (`Setup::check_from` refuses them first).
    pub fn setup(&mut self, tokens: &[Token]) -> Result<Vec<Resolved>, SessionError> {
        let mut out = Vec::new();
        for t in tokens {
            match t {
                Token::Register { alias, token_type, value } => {
                    if self.entries.contains_key(alias) {
                        return Err(SessionError::DUPLICATE_AUTH_TOKEN_ALIAS);
                    }
                    if self.used + size_of(value) <= self.max {
                        self.used += size_of(value);
                        self.entries.insert(*alias, (*token_type, value.clone()));
                    }
                    out.push((*token_type, value.clone()));
                }
                Token::UseValue { token_type, value } => out.push((*token_type, value.clone())),
                Token::Delete { .. } | Token::UseAlias { .. } => return Err(SessionError::PROTOCOL_VIOLATION),
            }
        }
        Ok(out)
    }

    /// The tokens of one request: registered, retired or resolved. A REGISTER that would
    /// pass the limit closes the session with AUTH_TOKEN_CACHE_OVERFLOW, and a duplicate
    /// alias with DUPLICATE_AUTH_TOKEN_ALIAS; a reference to an alias not registered is
    /// UNKNOWN_AUTH_TOKEN_ALIAS (§8.9). A REGISTER is kept even when the request it rides
    /// fails for another reason.
    pub fn resolve(&mut self, tokens: &[Token]) -> Result<Vec<Resolved>, SessionError> {
        let mut out = Vec::new();
        for t in tokens {
            match t {
                Token::Register { alias, token_type, value } => {
                    if self.entries.contains_key(alias) {
                        return Err(SessionError::DUPLICATE_AUTH_TOKEN_ALIAS);
                    }
                    if self.used + size_of(value) > self.max {
                        return Err(SessionError::AUTH_TOKEN_CACHE_OVERFLOW);
                    }
                    self.used += size_of(value);
                    self.entries.insert(*alias, (*token_type, value.clone()));
                    out.push((*token_type, value.clone()));
                }
                Token::Delete { alias } => match self.entries.remove(alias) {
                    Some((_, v)) => self.used -= size_of(&v),
                    None => return Err(SessionError::UNKNOWN_AUTH_TOKEN_ALIAS),
                },
                Token::UseAlias { alias } => match self.entries.get(alias) {
                    Some(r) => out.push(r.clone()),
                    None => return Err(SessionError::UNKNOWN_AUTH_TOKEN_ALIAS),
                },
                Token::UseValue { token_type, value } => out.push((*token_type, value.clone())),
            }
        }
        Ok(out)
    }
}

/// What a request asks to do.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Right {
    Publish,
    Subscribe,
}

struct GrantRule {
    prefix: TrackNamespace,
    token_sha256: [u8; 32],
    publish: bool,
    subscribe: bool,
    proof: bool,
    unlimited: bool,
    seek: bool,
}

struct AuthorityRule {
    key: [u8; 32],
    /// Empty: every room the authority grants.
    rooms: Vec<[u8; 32]>,
}

/// The relay's grants.
pub struct Authorizer {
    required: bool,
    grants: Vec<GrantRule>,
    authorities: Vec<AuthorityRule>,
}

pub fn sha256(b: &[u8]) -> [u8; 32] {
    crate::cert::sha256(b)
}

/// Whether `prefix` matches the first fields of `ns` (MOQT §2.4.2), a field `*` of
/// `prefix` matching any one field.
pub fn pattern_matches(prefix: &TrackNamespace, ns: &TrackNamespace) -> bool {
    prefix.0.len() <= ns.0.len() && prefix.0.iter().zip(&ns.0).all(|(p, n)| p.as_slice() == b"*" || p == n)
}

fn key32(what: &str, h: &str) -> Result<[u8; 32], String> {
    hex::decode(h).map_err(|e| format!("{what}: {e}"))?.try_into().map_err(|_| format!("{what} is not 32 bytes"))
}

fn unix_s() -> u64 {
    std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map_or(0, |d| d.as_secs())
}

/// Whether a room grant's head gives `right` over `ns` (WAI live-rooms §10.2, with the
/// room id `R` and the principal `P` as 64 lowercase hexadecimal digits): PUBLISH under
/// `[R, P, m]` for each media kind `m` it may publish, of `[R, P]` (its catalog track)
/// when it may publish any, and under `[R, P, events]` when its policy permits any event
/// kind; SUBSCRIBE and FETCH under `[R]` when it may subscribe. The relay's rights are per
/// namespace, so §10.2's track matches (the layers up to `max_layer`) are not narrowed
/// here.
pub fn room_grant_allows(g: &GrantHead, ns: &TrackNamespace, right: Right) -> bool {
    let r = hex::encode(g.room).into_bytes();
    let f = &ns.0;
    if f.first() != Some(&r) {
        return false;
    }
    match right {
        Right::Subscribe => g.subscribe,
        Right::Publish => {
            if f.len() < 2 || f[1] != hex::encode(g.principal).into_bytes() {
                return false;
            }
            match f.get(2).map(Vec::as_slice) {
                None => g.publish & 0b111 != 0,
                Some(b"events") => g.sends_events,
                Some(b"audio") => g.publishes(1),
                Some(b"video") => g.publishes(2),
                Some(b"screen") => g.publishes(3),
                Some(_) => false,
            }
        }
    }
}

impl Authorizer {
    pub fn from_config(c: &AuthConfig) -> Result<Self, String> {
        let mut grants = Vec::new();
        for g in &c.grants {
            let h = key32("a grant's token_sha256", &g.token_sha256)?;
            grants.push(GrantRule { prefix: TrackNamespace::from_strs(g.prefix.iter().map(String::as_str)), token_sha256: h, publish: g.publish, subscribe: g.subscribe, proof: g.proof, unlimited: g.unlimited, seek: g.seek });
        }
        let mut authorities = Vec::new();
        for a in &c.authorities {
            let key = key32("an authority's key", &a.key)?;
            ed25519_dalek::VerifyingKey::from_bytes(&key).map_err(|e| format!("an authority's key: {e}"))?;
            let rooms = a.rooms.iter().map(|r| key32("an authority's room", r)).collect::<Result<_, _>>()?;
            authorities.push(AuthorityRule { key, rooms });
        }
        Ok(Authorizer { required: c.required, grants, authorities })
    }

    pub fn open() -> Self {
        Authorizer { required: false, grants: Vec::new(), authorities: Vec::new() }
    }

    /// The head of `value` when it is a room grant one of the configured authorities
    /// signed, for a room that authority is honoured for.
    pub fn authority_grant(&self, value: &[u8]) -> Option<GrantHead> {
        if self.authorities.is_empty() {
            return None;
        }
        self.authorities.iter().find_map(|a| GrantHead::verified(value, &a.key).filter(|g| a.rooms.is_empty() || a.rooms.contains(&g.room)))
    }

    /// Whether `tokens` (the request's, and the session's from SETUP) give `right` over
    /// `ns`. A rule with `proof` counts only for a token whose SHA-256 `proven` holds: one
    /// the session's SETUP proved possession of. An authority's grant counts only when
    /// proven, and only within its window.
    pub fn allows(&self, ns: &TrackNamespace, right: Right, tokens: &[Resolved], proven: &HashSet<[u8; 32]>) -> bool {
        if !self.required {
            return true;
        }
        let now = unix_s();
        tokens.iter().any(|(ty, value)| {
            if *ty != 0 {
                return false;
            }
            let h = sha256(value);
            let by_rule = self.grants.iter().any(|g| {
                g.token_sha256 == h
                    && (!g.proof || proven.contains(&h))
                    && pattern_matches(&g.prefix, ns)
                    && match right {
                        Right::Publish => g.publish,
                        Right::Subscribe => g.subscribe,
                    }
            });
            by_rule || (proven.contains(&h) && self.authority_grant(value).is_some_and(|g| g.good_at(now) && room_grant_allows(&g, ns, right)))
        })
    }

    /// Whether `tokens` may seek back on a track under `ns`: fetch it, or fill it, from
    /// below its live edge (WAI live-rooms §10.6). With no authorization required, every
    /// session may. Otherwise a token whose rule gives subscribe and `seek` over `ns`, or an
    /// authority's grant, proved and within its window, that gives the track to subscribe
    /// (§10.2) and whose powers hold seek (§3.3, bit 4).
    pub fn may_seek(&self, ns: &TrackNamespace, tokens: &[Resolved], proven: &HashSet<[u8; 32]>) -> bool {
        if !self.required {
            return true;
        }
        let now = unix_s();
        tokens.iter().any(|(ty, value)| {
            if *ty != 0 {
                return false;
            }
            let h = sha256(value);
            let by_rule = self.grants.iter().any(|g| g.token_sha256 == h && (!g.proof || proven.contains(&h)) && pattern_matches(&g.prefix, ns) && g.subscribe && g.seek);
            by_rule || (proven.contains(&h) && self.authority_grant(value).is_some_and(|g| g.good_at(now) && g.seeks() && room_grant_allows(&g, ns, Right::Subscribe)))
        })
    }

    /// The heads of the authority-signed grants among `tokens` the session proved: whose
    /// powers a relay reads (a moderator's, WAI live-rooms §3.3).
    pub fn proven_grants(&self, tokens: &[Resolved], proven: &HashSet<[u8; 32]>) -> Vec<GrantHead> {
        tokens.iter().filter(|(ty, v)| *ty == 0 && proven.contains(&sha256(v))).filter_map(|(_, v)| self.authority_grant(v)).collect()
    }

    /// Whether `signer` is a room key of `room` (WAI live-rooms §2): an authority the relay is
    /// configured to honour for that room. A rating claim it signs about the room or a track
    /// of it is the room's (age-assurance §3.6).
    pub fn is_room_key(&self, signer: &[u8; 32], room: &[u8; 32]) -> bool {
        self.authorities.iter().any(|a| a.key == *signer && (a.rooms.is_empty() || a.rooms.contains(room)))
    }

    /// The keys of the configured room authorities.
    pub fn authority_keys(&self) -> Vec<[u8; 32]> {
        self.authorities.iter().map(|a| a.key).collect()
    }

    /// Whether the session's SETUP tokens name a rule that exempts it from the session and
    /// ingress limits.
    pub fn unlimited(&self, tokens: &[Resolved], proven: &HashSet<[u8; 32]>) -> bool {
        tokens.iter().any(|(ty, value)| {
            let h = sha256(value);
            *ty == 0 && self.grants.iter().any(|g| g.unlimited && g.token_sha256 == h && (!g.proof || proven.contains(&h)))
        })
    }

    /// Whether a token of type 0 with this value is a room grant honoured only beside a
    /// proof of possession: one a `proof` rule names, or one a configured authority signed.
    pub fn requires_proof(&self, value: &[u8]) -> bool {
        let h = sha256(value);
        self.required && (self.grants.iter().any(|g| g.proof && g.token_sha256 == h) || self.authority_grant(value).is_some())
    }
}

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

    #[test]
    fn the_cache_follows_section_8_9() {
        let mut c = TokenCache::new(16 + 4 + 16 + 3);
        let reg = |a, v: &[u8]| Token::Register { alias: a, token_type: 0, value: v.to_vec() };
        assert_eq!(c.resolve(&[reg(1, b"abcd")]).unwrap(), vec![(0, b"abcd".to_vec())]);
        assert_eq!(c.resolve(&[reg(1, b"x")]), Err(SessionError::DUPLICATE_AUTH_TOKEN_ALIAS));
        assert_eq!(c.resolve(&[reg(2, b"wxyz")]), Err(SessionError::AUTH_TOKEN_CACHE_OVERFLOW));
        assert_eq!(c.resolve(&[reg(2, b"xyz")]).unwrap().len(), 1);
        assert_eq!(c.resolve(&[Token::UseAlias { alias: 1 }]).unwrap(), vec![(0, b"abcd".to_vec())]);
        assert_eq!(c.resolve(&[Token::UseAlias { alias: 9 }]), Err(SessionError::UNKNOWN_AUTH_TOKEN_ALIAS));
        c.resolve(&[Token::Delete { alias: 1 }]).unwrap();
        assert_eq!(c.used(), 16 + 3);
        assert_eq!(c.resolve(&[Token::Delete { alias: 1 }]), Err(SessionError::UNKNOWN_AUTH_TOKEN_ALIAS));
        // In SETUP a registration past the limit reads as USE_VALUE.
        let mut s = TokenCache::new(0);
        assert_eq!(s.setup(&[reg(5, b"big")]).unwrap(), vec![(0, b"big".to_vec())]);
        assert_eq!(s.resolve(&[Token::UseAlias { alias: 5 }]), Err(SessionError::UNKNOWN_AUTH_TOKEN_ALIAS));
    }

    #[test]
    fn grants_cover_a_prefix_and_a_right() {
        let cfg = AuthConfig {
            required: true,
            authorities: vec![],
            grants: vec![
                crate::config::Grant { prefix: vec!["live".into()], token_sha256: hex::encode(sha256(b"t")), publish: false, subscribe: true, proof: false, unlimited: true, seek: false },
                crate::config::Grant { prefix: vec!["live".into()], token_sha256: hex::encode(sha256(b"g")), publish: true, subscribe: false, proof: true, unlimited: false, seek: false },
            ],
        };
        let a = Authorizer::from_config(&cfg).unwrap();
        let ns = TrackNamespace::from_strs(["live", "a"]);
        let none = HashSet::new();
        assert!(a.allows(&ns, Right::Subscribe, &[(0, b"t".to_vec())], &none));
        assert!(!a.allows(&ns, Right::Publish, &[(0, b"t".to_vec())], &none));
        assert!(!a.allows(&ns, Right::Subscribe, &[(0, b"u".to_vec())], &none));
        assert!(!a.allows(&TrackNamespace::from_strs(["other"]), Right::Subscribe, &[(0, b"t".to_vec())], &none));
        assert!(Authorizer::open().allows(&ns, Right::Publish, &[], &none));
        // A grant whose rule asks for a proof gives nothing until its session proved it.
        assert!(a.requires_proof(b"g") && !a.requires_proof(b"t"));
        assert!(!a.allows(&ns, Right::Publish, &[(0, b"g".to_vec())], &none));
        assert!(a.allows(&ns, Right::Publish, &[(0, b"g".to_vec())], &HashSet::from([sha256(b"g")])));
        // Exempt from the limits: a session presenting the token whose rule says so.
        assert!(a.unlimited(&[(0, b"t".to_vec())], &none) && !a.unlimited(&[(0, b"g".to_vec())], &none) && !a.unlimited(&[(0, b"u".to_vec())], &none));
    }

    /// A grant laid out as WAI live-rooms §4.1 lays it out, signed by `authority`.
    pub(crate) fn grant(authority: &ed25519_dalek::SigningKey, room: [u8; 32], principal: [u8; 32], nbf: u64, exp: u64, publish: u8, kinds: bool) -> Vec<u8> {
        use ed25519_dalek::Signer;
        let mut g = b"WRG1\x01".to_vec();
        g.extend_from_slice(&room);
        g.extend_from_slice(&principal);
        g.extend_from_slice(&[0; 40]); // prev, effective
        g.extend_from_slice(&nbf.to_le_bytes());
        g.extend_from_slice(&exp.to_le_bytes());
        g.extend_from_slice(&[0, publish, 1, 0, 0, 0]);
        g.extend_from_slice(&[0; 32]);
        g.extend_from_slice(b"WEP1\x01");
        let mut k = [0u8; 32];
        if kinds {
            k[0] = 1 << 3;
        }
        g.extend_from_slice(&k);
        g.extend_from_slice(&[0; 28]);
        g.extend_from_slice(&[0xff; 16]);
        g.extend_from_slice(&[0; 3]);
        g.extend_from_slice(&[0, 0, 0, 0]);
        let mut m = joule_moq::grant_proof::GRANT_DOMAIN.to_vec();
        m.extend_from_slice(&g);
        g.extend_from_slice(&authority.sign(&m).to_bytes());
        g
    }

    #[test]
    fn a_field_star_matches_any_field() {
        let p = TrackNamespace::from_strs(["*", "k"]);
        assert!(pattern_matches(&p, &TrackNamespace::from_strs(["r1", "k", "events"])));
        assert!(pattern_matches(&p, &TrackNamespace::from_strs(["r2", "k"])));
        assert!(!pattern_matches(&p, &TrackNamespace::from_strs(["r1", "j"])));
        assert!(!pattern_matches(&p, &TrackNamespace::from_strs(["r1"])));
        assert!(pattern_matches(&TrackNamespace::from_strs(Vec::<&str>::new()), &TrackNamespace::from_strs(["any"])));
    }

    #[test]
    fn an_authoritys_grant_maps_onto_rights_beside_its_proof() {
        let authority = ed25519_dalek::SigningKey::from_bytes(&[7; 32]);
        let stranger = ed25519_dalek::SigningKey::from_bytes(&[8; 32]);
        let (room, p) = ([0x11; 32], [0x22; 32]);
        let cfg = |rooms: Vec<String>| AuthConfig {
            required: true,
            grants: vec![],
            authorities: vec![crate::config::Authority { key: hex::encode(authority.verifying_key().to_bytes()), rooms }],
        };
        let a = Authorizer::from_config(&cfg(vec![])).unwrap();
        let now = unix_s();
        let g = grant(&authority, room, p, now - 60, now + 600, 0b011, true);
        let tok = vec![(0u64, g.clone())];
        let proven = HashSet::from([sha256(&g)]);
        let (r, ph) = (hex::encode(room), hex::encode(p));
        let ns = |f: &[&str]| TrackNamespace::from_strs(f.iter().copied());
        assert!(a.requires_proof(&g));
        // Without its proof, nothing.
        assert!(!a.allows(&ns(&[&r, &ph, "video"]), Right::Publish, &tok, &HashSet::new()));
        // With it: its own media, catalog and events; every track of the room to subscribe.
        for f in [vec![r.as_str(), ph.as_str(), "video"], vec![&r, &ph, "audio"], vec![&r, &ph], vec![&r, &ph, "events"]] {
            assert!(a.allows(&ns(&f), Right::Publish, &tok, &proven), "{f:?}");
        }
        assert!(!a.allows(&ns(&[&r, &ph, "screen"]), Right::Publish, &tok, &proven));
        assert!(!a.allows(&ns(&[&r, &"33".repeat(32), "video"]), Right::Publish, &tok, &proven), "another principal's");
        assert!(!a.allows(&ns(&[&"44".repeat(32), &ph, "video"]), Right::Publish, &tok, &proven), "another room");
        assert!(a.allows(&ns(&[&r, &"33".repeat(32), "video"]), Right::Subscribe, &tok, &proven));
        assert!(!a.allows(&ns(&[&"44".repeat(32)]), Right::Subscribe, &tok, &proven));
        // Outside its window, signed by another key, or for a room the rule does not name.
        let old = grant(&authority, room, p, now - 600, now - 60, 0b011, true);
        assert!(!a.allows(&ns(&[&r]), Right::Subscribe, &[(0, old.clone())], &HashSet::from([sha256(&old)])));
        let forged = grant(&stranger, room, p, now - 60, now + 600, 0b011, true);
        assert!(!a.requires_proof(&forged));
        assert!(!a.allows(&ns(&[&r]), Right::Subscribe, &[(0, forged.clone())], &HashSet::from([sha256(&forged)])));
        let narrow = Authorizer::from_config(&cfg(vec!["55".repeat(32)])).unwrap();
        assert!(!narrow.allows(&ns(&[&r]), Right::Subscribe, &tok, &proven) && !narrow.requires_proof(&g));
        // A viewer's grant: no media, no catalog, no events; the room to subscribe.
        let v = grant(&authority, room, p, now - 60, now + 600, 0, false);
        let vt = vec![(0u64, v.clone())];
        let vp = HashSet::from([sha256(&v)]);
        assert!(!a.allows(&ns(&[&r, &ph]), Right::Publish, &vt, &vp) && !a.allows(&ns(&[&r, &ph, "events"]), Right::Publish, &vt, &vp));
        assert!(a.allows(&ns(&[&r]), Right::Subscribe, &vt, &vp));
    }
}