world-id-primitives 0.14.0

use crate::{
    FieldElement, PrimitiveError,
    oprf::{OprfPrefix, OprfPrefixedFieldElement as _},
    poseidon::{self, ds},
};
use embed_doc_image::embed_doc_image;
use ruint::aliases::U256;
use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Error as _};

/// An identifier for a session (can be re-used).
///
/// A session allows RPs to ensure that it's still the same World ID
/// interacting with them across multiple interactions.
///
/// A `SessionId` is obtained after creating an initial session.
///
/// See the diagram below on how Session Proofs work, the [`SessionId`] and the `r` seed
/// ![Session Proofs Diagram][session-proofs.png]
///
/// Note that the `action` stored here is unrelated to the randomized action used
/// internally by [`SessionNullifier`]s — that randomized action exists only to ensure
/// the circuit's nullifier output is unique per Session Proof.
#[embed_doc_image("session-proofs.png", "assets/session-proofs.png")]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SessionId {
    /// The actual commitment being verified in the ZK-circuit.
    ///
    /// It is computed as H(DS_C || leaf_index || session_id_r_seed), see
    /// `signal computed_id_commitment` in `oprf_nullifier.circom`.
    pub commitment: FieldElement,
    /// A random seed generated by the authenticator in the initial Uniqueness Proof.
    ///
    /// This seed is the input to the OPRF Query to derive `session_id_r_seed` (`r`). It
    /// is part of the `session_id` so the RP can provide it when requesting a Session Proof.
    ///
    /// # Important: Prefix
    /// To ensure there are no collisions between the generated `r`s and the nullifiers
    /// for Uniqueness Proofs (as they use the same OPRF Key and query structure), the
    /// `oprf_seed`s, which are plugged as `action` in the Query Proof (see `QueryProofCircuitInput` in `world-id-proof`),
    /// MUST be prefixed with an explicit byte of `0x01`. All other actions have a `0x00` byte prefix. This
    /// collision avoidance is important because it ensures that any requests for nullifiers meant
    /// for Uniqueness Proofs are always signed by the RP (otherwise, an RP signature for a Session Proof
    /// could be used for requesting computation of _any_ nullifier).
    ///
    /// # Re-derivation
    ///
    /// The Authenticator can deterministically re-derive `r` from the OPRF nodes without
    /// needing to cache `r` locally as:
    /// ```text
    /// r = OPRF(pk_rpId, DS_C || leafIndex || oprf_seed)
    /// ```
    pub oprf_seed: FieldElement,
}

impl SessionId {
    const JSON_PREFIX: &str = "session_";

    /// Creates a new session id. Most uses should default to `from_r_seed` instead.
    ///
    /// # Errors
    /// If the provided `oprf_seed` is not prefixed properly.
    pub fn new(commitment: FieldElement, oprf_seed: FieldElement) -> Result<Self, PrimitiveError> {
        // OPRF Seeds must always start with a byte of `0x01`. See [`Self::oprf_seed`]
        // for details. Panic is acceptable as `oprf_seed` generation should
        // generally be done with `Self::from_r_seed`
        if !oprf_seed.has_prefix(OprfPrefix::SessionOprfSeed) {
            return Err(PrimitiveError::InvalidInput {
                attribute: "session_id".to_string(),
                reason: "inner oprf_seed is not valid".to_string(),
            });
        }
        Ok(Self {
            commitment,
            oprf_seed,
        })
    }

    /// Initializes a `SessionId` from the OPRF-output seed (`r`), and the `oprf_seed`
    /// used as input for the OPRF computation.
    ///
    /// This matches the logic in `oprf_nullifier.circom` for computing the `commitment` from the OPRF seed.
    ///
    /// # Seed (`session_id_r_seed`)
    /// - The seed MUST be computationally indistinguishable from random,
    ///   i.e. uniformly distributed because it uses OPRF.
    /// - When computed, the OPRF nodes will use the same `oprfKeyId` for the RP, with a different domain separator.
    /// - Requesting this seed requires a properly signed request from the RP and a complete query proof.
    /// - The seed generation is based on a randomly generated seed used as an "action" in a Query Proof. Note
    ///   this `action` is different than the randomized action used internally by [`SessionNullifier`]s.
    pub fn from_r_seed(
        leaf_index: u64,
        session_id_r_seed: FieldElement,
        oprf_seed: FieldElement,
    ) -> Result<Self, PrimitiveError> {
        if !oprf_seed.has_prefix(OprfPrefix::SessionOprfSeed) {
            return Err(PrimitiveError::InvalidInput {
                attribute: "session_id".to_string(),
                reason: "inner oprf_seed is not valid".to_string(),
            });
        }

        let commitment = poseidon::hash(
            ds::SESSION_COMMITMENT,
            [leaf_index.into(), session_id_r_seed],
        );
        Ok(Self {
            commitment,
            oprf_seed,
        })
    }

    /// Generates a new [`Self::oprf_seed`] to initialize a new Session.
    pub fn generate_oprf_seed<R: rand::CryptoRng + rand::RngCore>(rng: &mut R) -> FieldElement {
        FieldElement::random_with_prefix(rng, OprfPrefix::SessionOprfSeed)
    }

    /// Verifies that `session_id_r_seed` re-derives to this session id's [`Self::commitment`].
    ///
    /// Checks an `r` seed obtained elsewhere (e.g. an Authenticator's cache) without
    /// re-deriving it via the OPRF nodes.
    ///
    /// # Errors
    /// - If this session id's [`Self::oprf_seed`] is not valid.
    /// - If the derived commitment does not match [`Self::commitment`].
    pub fn verify_commitment(
        &self,
        leaf_index: u64,
        session_id_r_seed: FieldElement,
    ) -> Result<(), PrimitiveError> {
        let computed = Self::from_r_seed(leaf_index, session_id_r_seed, self.oprf_seed)?;
        if computed.commitment != self.commitment {
            return Err(PrimitiveError::SessionIdCommitmentMismatch);
        }
        Ok(())
    }

    /// Returns the 64-byte big-endian representation (2 x 32-byte field elements).
    #[must_use]
    pub fn to_compressed_bytes(&self) -> [u8; 64] {
        let mut bytes = [0u8; 64];
        bytes[..32].copy_from_slice(&self.commitment.to_be_bytes());
        bytes[32..].copy_from_slice(&self.oprf_seed.to_be_bytes());
        bytes
    }

    /// Constructs from compressed bytes (must be exactly 64 bytes).
    ///
    /// # Errors
    /// Returns an error if the input is not exactly 64 bytes or if values are not valid field elements.
    pub fn from_compressed_bytes(bytes: &[u8]) -> Result<Self, String> {
        if bytes.len() != 64 {
            return Err(format!(
                "Invalid length: expected 64 bytes, got {}",
                bytes.len()
            ));
        }

        let commitment = FieldElement::from_be_bytes(bytes[..32].try_into().unwrap())
            .map_err(|e| format!("invalid commitment: {e}"))?;
        let oprf_seed = FieldElement::from_be_bytes(bytes[32..].try_into().unwrap())
            .map_err(|e| format!("invalid oprf_seed: {e}"))?;

        if bytes[32] != OprfPrefix::SessionOprfSeed as u8 {
            return Err("invalid prefix for oprf_seed".to_string());
        }

        Ok(Self {
            commitment,
            oprf_seed,
        })
    }
}

impl Default for SessionId {
    fn default() -> Self {
        let mut oprf_seed = [0u8; 32];
        oprf_seed[0] = OprfPrefix::SessionOprfSeed as u8;
        let oprf_seed = U256::from_be_bytes(oprf_seed)
            .try_into()
            .expect("always fits in the field");
        Self {
            commitment: FieldElement::ZERO,
            oprf_seed,
        }
    }
}

impl Serialize for SessionId {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let bytes = self.to_compressed_bytes();
        if serializer.is_human_readable() {
            // JSON: prefixed hex-encoded compressed bytes for explicit typing.
            serializer.serialize_str(&format!("{}{}", Self::JSON_PREFIX, hex::encode(bytes)))
        } else {
            // Binary: compressed bytes
            serializer.serialize_bytes(&bytes)
        }
    }
}

impl<'de> Deserialize<'de> for SessionId {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let bytes = if deserializer.is_human_readable() {
            let value = String::deserialize(deserializer)?;
            let hex_str = value.strip_prefix(Self::JSON_PREFIX).ok_or_else(|| {
                D::Error::custom(format!(
                    "session id must start with '{}'",
                    Self::JSON_PREFIX
                ))
            })?;
            hex::decode(hex_str).map_err(D::Error::custom)?
        } else {
            Vec::deserialize(deserializer)?
        };

        Self::from_compressed_bytes(&bytes).map_err(D::Error::custom)
    }
}

/// How a proof request refers to a session.
///
/// Wire encoding (the request's `session_id` field): absent or `null` → [`Self::None`],
/// `"create"` → [`Self::Create`], a `"session_"`-prefixed id → [`Self::Existing`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub enum SessionRef {
    /// No session involvement.
    #[default]
    None,
    /// Mint a fresh session. For [`crate::ProofType::Session`] this proves the new
    /// session in the same response; for [`crate::ProofType::Uniqueness`] this
    /// returns a uniqueness proof committed to the newly minted session.
    Create,
    /// Refer to an existing session. Only valid for [`crate::ProofType::Session`].
    Existing(SessionId),
}

impl SessionRef {
    const CREATE_TOKEN: &str = "create";

    /// Returns true if the request involves no session.
    #[must_use]
    pub const fn is_none(&self) -> bool {
        matches!(self, Self::None)
    }

    /// Returns true if the request asks to mint a fresh session.
    #[must_use]
    pub const fn is_create(&self) -> bool {
        matches!(self, Self::Create)
    }

    /// Returns the referenced existing session id, if any.
    #[must_use]
    pub const fn existing(&self) -> Option<SessionId> {
        match self {
            Self::Existing(id) => Some(*id),
            _ => None,
        }
    }
}

impl From<SessionId> for SessionRef {
    fn from(id: SessionId) -> Self {
        Self::Existing(id)
    }
}

impl Serialize for SessionRef {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        match self {
            Self::None => serializer.serialize_none(),
            Self::Create => {
                if serializer.is_human_readable() {
                    serializer.serialize_str(Self::CREATE_TOKEN)
                } else {
                    // Binary: 6-byte token, cannot collide with the 64-byte `SessionId` encoding
                    serializer.serialize_bytes(Self::CREATE_TOKEN.as_bytes())
                }
            }
            Self::Existing(id) => id.serialize(serializer),
        }
    }
}

impl<'de> Deserialize<'de> for SessionRef {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        struct SessionRefVisitor;

        impl<'de> serde::de::Visitor<'de> for SessionRefVisitor {
            type Value = SessionRef;

            fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
                write!(
                    formatter,
                    "null, \"{}\", or a '{}'-prefixed session id",
                    SessionRef::CREATE_TOKEN,
                    SessionId::JSON_PREFIX
                )
            }

            fn visit_none<E: serde::de::Error>(self) -> Result<Self::Value, E> {
                Ok(SessionRef::None)
            }

            fn visit_unit<E: serde::de::Error>(self) -> Result<Self::Value, E> {
                Ok(SessionRef::None)
            }

            fn visit_some<D2>(self, deserializer: D2) -> Result<Self::Value, D2::Error>
            where
                D2: Deserializer<'de>,
            {
                if deserializer.is_human_readable() {
                    let value = String::deserialize(deserializer)?;
                    if value == SessionRef::CREATE_TOKEN {
                        return Ok(SessionRef::Create);
                    }
                    let hex_str = value.strip_prefix(SessionId::JSON_PREFIX).ok_or_else(|| {
                        D2::Error::custom(format!(
                            "session_id must be \"{}\" or start with '{}'",
                            SessionRef::CREATE_TOKEN,
                            SessionId::JSON_PREFIX
                        ))
                    })?;
                    let bytes = hex::decode(hex_str).map_err(D2::Error::custom)?;
                    SessionId::from_compressed_bytes(&bytes)
                        .map(SessionRef::Existing)
                        .map_err(D2::Error::custom)
                } else {
                    let bytes = Vec::<u8>::deserialize(deserializer)?;
                    if bytes == SessionRef::CREATE_TOKEN.as_bytes() {
                        return Ok(SessionRef::Create);
                    }
                    SessionId::from_compressed_bytes(&bytes)
                        .map(SessionRef::Existing)
                        .map_err(D2::Error::custom)
                }
            }
        }

        deserializer.deserialize_option(SessionRefVisitor)
    }
}

/// A session nullifier for World ID Session proofs. It is analogous to a request nonce,
/// it **does NOT guarantee uniqueness of a World ID** as a `Nullifier` does.
///
/// This type is intended to be opaque for RPs. For an RP context, they should only
/// be concerned of this needing to be passthrough to the `verifySession()` contract function.
///
/// This type exists as an adaptation to be able to use the same ZK-circuit for
/// both Uniqueness Proofs and Session Proofs, and it encompasses:
/// - the nullifier used as the proof output.
/// - a random action bound to the same proof.
///
/// The `WorldIDVerifier.sol` contract expects this as a `uint256[2]` array
/// use `as_ethereum_representation()` for conversion.
///
/// # Future
///
/// Note the session nullifier exists **only** to support the same ZK-circuit than for Uniqueness Proofs; as
/// World ID evolves to a different proving system which won't require circuit precompiles, a new circuit MUST
/// be created which does not generate a nullifier at all, and the input randomized action will not be required either.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SessionNullifier {
    /// The nullifier value for this proof.
    nullifier: FieldElement,
    /// The random action value bound to this session proof.
    action: FieldElement,
}

impl SessionNullifier {
    const JSON_PREFIX: &str = "snil_";

    /// Creates a new session nullifier.
    pub fn new(nullifier: FieldElement, action: FieldElement) -> Result<Self, PrimitiveError> {
        if !action.has_prefix(OprfPrefix::SessionAction) {
            return Err(PrimitiveError::InvalidInput {
                attribute: "session_nullifier".to_string(),
                reason: "inner action is not valid".to_string(),
            });
        }
        Ok(Self { nullifier, action })
    }

    /// Returns the nullifier value.
    #[must_use]
    pub const fn nullifier(&self) -> FieldElement {
        self.nullifier
    }

    /// Returns the action value.
    #[must_use]
    pub const fn action(&self) -> FieldElement {
        self.action
    }

    /// Returns the session nullifier as an Ethereum-compatible array for `verifySession()`.
    ///
    /// Format: `[nullifier, action]` matching the contract's `uint256[2] sessionNullifier`.
    #[must_use]
    pub fn as_ethereum_representation(&self) -> [U256; 2] {
        [self.nullifier.into(), self.action.into()]
    }

    /// Creates a session nullifier from an Ethereum representation.
    ///
    /// # Errors
    /// Returns an error if the U256 values are not valid field elements.
    pub fn from_ethereum_representation(value: [U256; 2]) -> Result<Self, String> {
        let nullifier =
            FieldElement::try_from(value[0]).map_err(|e| format!("invalid nullifier: {e}"))?;
        let action =
            FieldElement::try_from(value[1]).map_err(|e| format!("invalid action: {e}"))?;

        if !action.has_prefix(OprfPrefix::SessionAction) {
            return Err("inner action is not valid".to_string());
        }
        Ok(Self { nullifier, action })
    }

    /// Returns the 64-byte big-endian representation (2 x 32-byte field elements).
    #[must_use]
    pub fn to_compressed_bytes(&self) -> [u8; 64] {
        let mut bytes = [0u8; 64];
        bytes[..32].copy_from_slice(&self.nullifier.to_be_bytes());
        bytes[32..].copy_from_slice(&self.action.to_be_bytes());
        bytes
    }

    /// Constructs from compressed bytes (must be exactly 64 bytes).
    ///
    /// # Errors
    /// Returns an error if the input is not exactly 64 bytes or if values are not valid field elements.
    pub fn from_compressed_bytes(bytes: &[u8]) -> Result<Self, String> {
        if bytes.len() != 64 {
            return Err(format!(
                "Invalid length: expected 64 bytes, got {}",
                bytes.len()
            ));
        }

        let nullifier = FieldElement::from_be_bytes(bytes[..32].try_into().unwrap())
            .map_err(|e| format!("invalid nullifier: {e}"))?;
        let action = FieldElement::from_be_bytes(bytes[32..].try_into().unwrap())
            .map_err(|e| format!("invalid action: {e}"))?;

        if bytes[32] != OprfPrefix::SessionAction as u8 {
            return Err("invalid action. missing expected prefix.".to_string());
        }

        Ok(Self { nullifier, action })
    }
}

impl Default for SessionNullifier {
    fn default() -> Self {
        let mut action = [0u8; 32];
        action[0] = OprfPrefix::SessionAction as u8;
        let action = U256::from_be_bytes(action)
            .try_into()
            .expect("always fits in the field");
        Self {
            nullifier: FieldElement::ZERO,
            action,
        }
    }
}

impl Serialize for SessionNullifier {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let bytes = self.to_compressed_bytes();
        if serializer.is_human_readable() {
            // JSON: prefixed hex-encoded compressed bytes for explicit typing.
            serializer.serialize_str(&format!("{}{}", Self::JSON_PREFIX, hex::encode(bytes)))
        } else {
            // Binary: compressed bytes
            serializer.serialize_bytes(&bytes)
        }
    }
}

impl<'de> Deserialize<'de> for SessionNullifier {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let bytes = if deserializer.is_human_readable() {
            let value = String::deserialize(deserializer)?;
            let hex_str = value.strip_prefix(Self::JSON_PREFIX).ok_or_else(|| {
                D::Error::custom(format!(
                    "session nullifier must start with '{}'",
                    Self::JSON_PREFIX
                ))
            })?;
            hex::decode(hex_str).map_err(D::Error::custom)?
        } else {
            Vec::deserialize(deserializer)?
        };

        Self::from_compressed_bytes(&bytes).map_err(D::Error::custom)
    }
}

impl From<SessionNullifier> for [U256; 2] {
    fn from(value: SessionNullifier) -> Self {
        value.as_ethereum_representation()
    }
}

#[cfg(test)]
mod session_id_tests {
    use super::*;
    use ruint::uint;

    fn test_field_element(value: u64) -> FieldElement {
        FieldElement::from(value)
    }

    /// Creates an oprf_seed with the right prefix
    fn test_oprf_seed(value: u64) -> FieldElement {
        // set the first byte to 0x01; no need to clear the first bits as the input is u64
        let n = U256::from(value)
            | uint!(0x0100000000000000000000000000000000000000000000000000000000000000_U256);
        FieldElement::try_from(n).expect("test value fits in field")
    }

    #[test]
    fn test_new_and_accessors() {
        let commitment = test_field_element(1001);
        let seed = test_oprf_seed(42);
        let id = SessionId::new(commitment, seed).unwrap();

        assert_eq!(id.commitment, commitment);
        assert_eq!(id.oprf_seed, seed);
    }

    #[test]
    fn test_default() {
        let id = SessionId::default();
        assert_eq!(id.commitment, FieldElement::ZERO);
        assert_eq!(
            id.oprf_seed,
            uint!(0x0100000000000000000000000000000000000000000000000000000000000000_U256)
                .try_into()
                .unwrap()
        );
    }

    #[test]
    fn test_bytes_roundtrip() {
        let id = SessionId::new(test_field_element(1001), test_oprf_seed(42)).unwrap();
        let bytes = id.to_compressed_bytes();

        assert_eq!(bytes.len(), 64);

        let decoded = SessionId::from_compressed_bytes(&bytes).unwrap();
        assert_eq!(id, decoded);
    }

    #[test]
    fn test_bytes_use_field_element_encoding() {
        let id = SessionId::new(test_field_element(1001), test_oprf_seed(42)).unwrap();
        let bytes = id.to_compressed_bytes();

        let mut expected = [0u8; 64];
        expected[..32].copy_from_slice(&id.commitment.to_be_bytes());
        expected[32..].copy_from_slice(&id.oprf_seed.to_be_bytes());
        assert_eq!(bytes, expected);
    }

    #[test]
    fn test_invalid_bytes_length() {
        let too_short = vec![0u8; 63];
        let result = SessionId::from_compressed_bytes(&too_short);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("Invalid length"));

        let too_long = vec![0u8; 65];
        let result = SessionId::from_compressed_bytes(&too_long);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("Invalid length"));
    }

    #[test]
    fn test_from_compressed_bytes_rejects_wrong_oprf_seed_prefix() {
        let mut bytes = [0u8; 64];
        // Valid commitment (zero is a valid field element)
        // oprf_seed with wrong prefix: 0x00 instead of 0x01
        bytes[32] = 0x00;
        let result = SessionId::from_compressed_bytes(&bytes);
        assert!(result.is_err());
        assert!(
            result.unwrap_err().contains("invalid prefix"),
            "should reject oprf_seed without 0x01 prefix"
        );
    }

    #[test]
    fn test_json_roundtrip() {
        let id = SessionId::new(test_field_element(1001), test_oprf_seed(42)).unwrap();
        let json = serde_json::to_string(&id).unwrap();

        assert!(json.starts_with("\"session_"));
        assert!(json.ends_with('"'));

        let decoded: SessionId = serde_json::from_str(&json).unwrap();
        assert_eq!(id, decoded);
    }

    #[test]
    fn test_json_format() {
        let id = SessionId::new(test_field_element(1), test_oprf_seed(2)).unwrap();
        let json = serde_json::to_string(&id).unwrap();

        let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
        assert!(parsed.is_string());
        let value = parsed.as_str().unwrap();
        assert!(value.starts_with("session_"));
    }

    #[test]
    fn test_json_wrong_prefix_rejected() {
        let result = serde_json::from_str::<SessionId>("\"snil_00\"");
        assert!(result.is_err());
    }

    #[test]
    fn test_generates_random_oprf_seed() {
        let mut rng = rand::rngs::OsRng;

        let seed_1 = SessionId::generate_oprf_seed(&mut rng);
        let seed_2 = SessionId::generate_oprf_seed(&mut rng);

        assert_ne!(seed_1, seed_2);
    }

    #[test]
    fn test_from_r_seed_generated_seed_has_session_prefix() {
        let mut rng = rand::rngs::OsRng;

        for _ in 0..50 {
            let seed = SessionId::generate_oprf_seed(&mut rng);
            // Top byte must be exactly 0x01: bit 248 set, bits 249-255 clear
            assert_eq!(seed.to_u256() >> 248, U256::from(1));
        }
    }

    #[test]
    fn test_from_r_seed_commitment_snapshot() {
        let leaf_index = 42u64;
        let r_seed = test_field_element(123);
        let oprf_seed = test_oprf_seed(456);

        let session_id = SessionId::from_r_seed(leaf_index, r_seed, oprf_seed).unwrap();

        let expected = "0x1e7853ebd4fc9d9f0232fdcfae116023610bdf66a22e2700445d7a2e0e7e6152"
            .parse::<U256>()
            .unwrap();
        assert_eq!(
            session_id.commitment.to_u256(),
            expected,
            "commitment snapashot for session commitment changed"
        );
    }

    #[test]
    fn test_verify_commitment() {
        let leaf_index = 42u64;
        let r_seed = test_field_element(123);
        let session_id = SessionId::from_r_seed(leaf_index, r_seed, test_oprf_seed(456)).unwrap();

        session_id.verify_commitment(leaf_index, r_seed).unwrap();
        assert_eq!(
            session_id
                .verify_commitment(leaf_index, test_field_element(124))
                .unwrap_err(),
            PrimitiveError::SessionIdCommitmentMismatch
        );
    }
}

#[cfg(test)]
mod session_ref_tests {
    use super::*;
    use ruint::uint;

    fn test_session_id() -> SessionId {
        let oprf_seed = U256::from(42u64)
            | uint!(0x0100000000000000000000000000000000000000000000000000000000000000_U256);
        SessionId::new(
            FieldElement::from(1001u64),
            FieldElement::try_from(oprf_seed).expect("test value fits in field"),
        )
        .expect("valid session id")
    }

    #[test]
    fn test_default_is_none() {
        assert_eq!(SessionRef::default(), SessionRef::None);
        assert!(SessionRef::None.is_none());
        assert!(SessionRef::Create.is_create());
        assert_eq!(
            SessionRef::Existing(test_session_id()).existing(),
            Some(test_session_id())
        );
        assert_eq!(
            SessionRef::from(test_session_id()).existing(),
            Some(test_session_id())
        );
    }

    #[test]
    fn test_deserialize_create_token() {
        let parsed: SessionRef = serde_json::from_str("\"create\"").unwrap();
        assert_eq!(parsed, SessionRef::Create);
    }

    #[test]
    fn test_deserialize_existing_matches_session_id_parse() {
        let id = test_session_id();
        let json = serde_json::to_string(&id).unwrap();
        let parsed: SessionRef = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed, SessionRef::Existing(id));
    }

    #[test]
    fn test_deserialize_null_is_none() {
        let parsed: SessionRef = serde_json::from_str("null").unwrap();
        assert_eq!(parsed, SessionRef::None);
    }

    #[test]
    fn test_rejects_unknown_strings() {
        for input in ["\"Create\"", "\"creat\"", "\"snil_00\"", "\"\""] {
            let result = serde_json::from_str::<SessionRef>(input);
            let err = result.expect_err(input).to_string();
            assert!(
                err.contains("create") || err.contains("session_"),
                "error for {input} should name the accepted forms: {err}"
            );
        }
    }

    #[test]
    fn test_json_roundtrip_all_states() {
        let cases = [
            (SessionRef::None, "null"),
            (SessionRef::Create, "\"create\""),
        ];
        for (state, expected_json) in cases {
            let json = serde_json::to_string(&state).unwrap();
            assert_eq!(json, expected_json);
            let parsed: SessionRef = serde_json::from_str(&json).unwrap();
            assert_eq!(parsed, state);
        }

        let existing = SessionRef::Existing(test_session_id());
        let json = serde_json::to_string(&existing).unwrap();
        assert!(json.starts_with("\"session_"));
        let parsed: SessionRef = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed, existing);
    }

    #[test]
    fn test_cbor_roundtrip_all_states() {
        for state in [
            SessionRef::None,
            SessionRef::Create,
            SessionRef::Existing(test_session_id()),
        ] {
            let mut buffer = Vec::new();
            ciborium::into_writer(&state, &mut buffer).unwrap();
            let decoded: SessionRef = ciborium::from_reader(&buffer[..]).unwrap();
            assert_eq!(state, decoded);
        }
    }
}

#[cfg(test)]
mod session_nullifier_tests {
    use super::*;
    use ruint::uint;

    fn test_field_element(value: u64) -> FieldElement {
        FieldElement::from(value)
    }

    /// Creates an action with the required `0x02` prefix
    fn test_action(value: u64) -> FieldElement {
        let n = U256::from(value)
            | uint!(0x0200000000000000000000000000000000000000000000000000000000000000_U256);
        FieldElement::try_from(n).expect("test value fits in field")
    }

    #[test]
    fn test_new_and_accessors() {
        let nullifier = test_field_element(1001);
        let action = test_action(42);
        let session = SessionNullifier::new(nullifier, action).unwrap();

        assert_eq!(session.nullifier(), nullifier);
        assert_eq!(session.action(), action);
    }

    #[test]
    fn test_as_ethereum_representation() {
        let nullifier = test_field_element(100);
        let action = test_action(200);
        let session = SessionNullifier::new(nullifier, action).unwrap();

        let repr = session.as_ethereum_representation();
        assert_eq!(repr[0], U256::from(100));
        assert_eq!(repr[1], action.to_u256());
    }

    #[test]
    fn test_from_ethereum_representation() {
        let action = test_action(200);
        let repr = [U256::from(100), action.to_u256()];
        let session = SessionNullifier::from_ethereum_representation(repr).unwrap();

        assert_eq!(session.nullifier(), test_field_element(100));
        assert_eq!(session.action(), action);
    }

    #[test]
    fn test_json_roundtrip() {
        let session = SessionNullifier::new(test_field_element(1001), test_action(42)).unwrap();
        let json = serde_json::to_string(&session).unwrap();

        // Verify JSON uses the prefixed compact representation
        assert!(json.starts_with("\"snil_"));
        assert!(json.ends_with('"'));

        // Verify roundtrip
        let decoded: SessionNullifier = serde_json::from_str(&json).unwrap();
        assert_eq!(session, decoded);
    }

    #[test]
    fn test_json_format() {
        let session = SessionNullifier::new(test_field_element(1), test_action(2)).unwrap();
        let json = serde_json::to_string(&session).unwrap();

        // Should be a prefixed compact string
        let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
        assert!(parsed.is_string());
        let value = parsed.as_str().unwrap();
        assert!(value.starts_with("snil_"));
    }

    #[test]
    fn test_bytes_roundtrip() {
        let session = SessionNullifier::new(test_field_element(1001), test_action(42)).unwrap();
        let bytes = session.to_compressed_bytes();

        assert_eq!(bytes.len(), 64); // 32 + 32 bytes

        let decoded = SessionNullifier::from_compressed_bytes(&bytes).unwrap();
        assert_eq!(session, decoded);
    }

    #[test]
    fn test_bytes_use_field_element_encoding() {
        let session = SessionNullifier::new(test_field_element(1001), test_action(42)).unwrap();
        let bytes = session.to_compressed_bytes();

        let mut expected = [0u8; 64];
        expected[..32].copy_from_slice(&session.nullifier().to_be_bytes());
        expected[32..].copy_from_slice(&session.action().to_be_bytes());
        assert_eq!(bytes, expected);
    }

    #[test]
    fn test_invalid_bytes_length() {
        let too_short = vec![0u8; 63];
        let result = SessionNullifier::from_compressed_bytes(&too_short);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("Invalid length"));

        let too_long = vec![0u8; 65];
        let result = SessionNullifier::from_compressed_bytes(&too_long);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("Invalid length"));
    }

    #[test]
    fn test_default() {
        let session = SessionNullifier::default();
        assert_eq!(session.nullifier(), FieldElement::ZERO);
        let expected_action: FieldElement =
            uint!(0x0200000000000000000000000000000000000000000000000000000000000000_U256)
                .try_into()
                .unwrap();
        assert_eq!(session.action(), expected_action);
    }

    #[test]
    fn test_into_u256_array() {
        let action = test_action(200);
        let session = SessionNullifier::new(test_field_element(100), action).unwrap();
        let arr: [U256; 2] = session.into();

        assert_eq!(arr[0], U256::from(100));
        assert_eq!(arr[1], action.to_u256());
    }

    #[test]
    fn test_new_rejects_invalid_action_prefix() {
        let nullifier = test_field_element(1);
        let bad_action = test_field_element(42); // no 0x02 prefix
        let result = SessionNullifier::new(nullifier, bad_action);
        assert!(result.is_err());

        let err = result.unwrap_err();
        assert!(
            matches!(err, PrimitiveError::InvalidInput { .. }),
            "expected InvalidInput, got {err:?}"
        );
    }

    #[test]
    fn test_new_rejects_oprf_seed_prefix_as_action() {
        let nullifier = test_field_element(1);
        // 0x01 prefix (OprfSeed) is not valid for Action
        let oprf_prefixed = U256::from(42u64)
            | uint!(0x0100000000000000000000000000000000000000000000000000000000000000_U256);
        let bad_action = FieldElement::try_from(oprf_prefixed).unwrap();
        assert!(SessionNullifier::new(nullifier, bad_action).is_err());
    }

    #[test]
    fn test_from_ethereum_representation_rejects_invalid_action() {
        let repr = [U256::from(100), U256::from(200)]; // action has 0x00 prefix
        let result = SessionNullifier::from_ethereum_representation(repr);
        assert!(result.is_err());
        assert!(
            result.unwrap_err().contains("action"),
            "error should mention the action"
        );
    }

    #[test]
    fn test_from_compressed_bytes_rejects_invalid_action_prefix() {
        let mut bytes = [0u8; 64];
        // Valid nullifier (zero), but action with 0x00 prefix
        bytes[32] = 0x00;
        let result = SessionNullifier::from_compressed_bytes(&bytes);
        assert!(result.is_err());
        assert!(
            result.unwrap_err().contains("action"),
            "error should mention the action"
        );
    }

    #[test]
    fn test_json_rejects_invalid_action_prefix() {
        // Build JSON with a valid nullifier but an action lacking the 0x02 prefix
        let nullifier = test_field_element(1);
        let bad_action = test_field_element(2); // 0x00 prefix
        let mut bytes = [0u8; 64];
        bytes[..32].copy_from_slice(&nullifier.to_be_bytes());
        bytes[32..].copy_from_slice(&bad_action.to_be_bytes());
        let json = format!("\"snil_{}\"", hex::encode(bytes));

        let result = serde_json::from_str::<SessionNullifier>(&json);
        assert!(result.is_err());
    }
}