//! Canonical Peach Noise state machine.
//!
//! The browser and Android adapters are deliberately thin: all NK/KK
//! transcript processing, DH, AEAD, nonce handling, and transport rekeying
//! happen here through `snow`'s unmodified named patterns.

use snow::{
    Builder, Error, HandshakeState, TransportState,
    params::{CipherChoice, DHChoice, HashChoice, NoiseParams},
    resolvers::{CryptoResolver, DefaultResolver},
    types::{Cipher, Dh, Hash, Random},
};
use subtle::ConstantTimeEq;
use x25519_dalek::{PublicKey, StaticSecret};
use zeroize::{Zeroize, Zeroizing};

const NK: &str = "Noise_NK_25519_ChaChaPoly_BLAKE2s";
const KK: &str = "Noise_KK_25519_ChaChaPoly_BLAKE2s";
const HASH_LEN: usize = 32;
const TAG_LEN: usize = 16;
const MAX_NOISE_MESSAGE: usize = 65_535;

type CoreResult<T> = Result<T, String>;

enum SessionState {
    Handshake(Option<HandshakeState>),
    Transport(TransportState),
    Closed,
}

pub struct CoreNoiseSession {
    state: SessionState,
    handshake_hash: Option<[u8; HASH_LEN]>,
}

impl CoreNoiseSession {
    pub fn nk_initiator(remote_public_key: &[u8]) -> CoreResult<Self> {
        Self::nk_initiator_inner(remote_public_key, None)
    }

    #[doc(hidden)]
    pub fn nk_initiator_with_ephemeral_for_test(
        remote_public_key: &[u8],
        ephemeral_private_key: &[u8],
    ) -> CoreResult<Self> {
        Self::nk_initiator_inner(remote_public_key, Some(ephemeral_private_key))
    }

    fn nk_initiator_inner(
        remote_public_key: &[u8],
        ephemeral_private_key: Option<&[u8]>,
    ) -> CoreResult<Self> {
        require_key(remote_public_key, "responder public key")?;
        let params: NoiseParams = NK.parse().map_err(display_error)?;
        let mut builder = core_builder(params)
            .remote_public_key(remote_public_key)
            .map_err(display_error)?;
        if let Some(ephemeral) = ephemeral_private_key {
            require_key(ephemeral, "initiator ephemeral private key")?;
            builder = builder.fixed_ephemeral_key_for_testing_only(ephemeral);
        }
        let handshake = builder.build_initiator().map_err(display_error)?;
        Ok(Self::handshake(handshake))
    }

    pub fn nk_responder(local_private_key: &[u8], local_public_key: &[u8]) -> CoreResult<Self> {
        Self::nk_responder_inner(local_private_key, local_public_key, None)
    }

    #[doc(hidden)]
    pub fn nk_responder_with_ephemeral_for_test(
        local_private_key: &[u8],
        local_public_key: &[u8],
        ephemeral_private_key: &[u8],
    ) -> CoreResult<Self> {
        Self::nk_responder_inner(
            local_private_key,
            local_public_key,
            Some(ephemeral_private_key),
        )
    }

    fn nk_responder_inner(
        local_private_key: &[u8],
        local_public_key: &[u8],
        ephemeral_private_key: Option<&[u8]>,
    ) -> CoreResult<Self> {
        require_matching_keypair(local_private_key, local_public_key)?;
        let private = Zeroizing::new(local_private_key.to_vec());
        let params: NoiseParams = NK.parse().map_err(display_error)?;
        let mut builder = core_builder(params)
            .local_private_key(private.as_slice())
            .map_err(display_error)?;
        if let Some(ephemeral) = ephemeral_private_key {
            require_key(ephemeral, "responder ephemeral private key")?;
            builder = builder.fixed_ephemeral_key_for_testing_only(ephemeral);
        }
        let handshake = builder.build_responder().map_err(display_error)?;
        Ok(Self::handshake(handshake))
    }

    pub fn kk_initiator(
        local_private_key: &[u8],
        local_public_key: &[u8],
        remote_public_key: &[u8],
    ) -> CoreResult<Self> {
        Self::kk_initiator_inner(local_private_key, local_public_key, remote_public_key, None)
    }

    #[doc(hidden)]
    pub fn kk_initiator_with_ephemeral_for_test(
        local_private_key: &[u8],
        local_public_key: &[u8],
        remote_public_key: &[u8],
        ephemeral_private_key: &[u8],
    ) -> CoreResult<Self> {
        Self::kk_initiator_inner(
            local_private_key,
            local_public_key,
            remote_public_key,
            Some(ephemeral_private_key),
        )
    }

    fn kk_initiator_inner(
        local_private_key: &[u8],
        local_public_key: &[u8],
        remote_public_key: &[u8],
        ephemeral_private_key: Option<&[u8]>,
    ) -> CoreResult<Self> {
        require_matching_keypair(local_private_key, local_public_key)?;
        require_key(remote_public_key, "responder public key")?;
        let private = Zeroizing::new(local_private_key.to_vec());
        let params: NoiseParams = KK.parse().map_err(display_error)?;
        let mut builder = core_builder(params)
            .local_private_key(private.as_slice())
            .map_err(display_error)?
            .remote_public_key(remote_public_key)
            .map_err(display_error)?;
        if let Some(ephemeral) = ephemeral_private_key {
            require_key(ephemeral, "initiator ephemeral private key")?;
            builder = builder.fixed_ephemeral_key_for_testing_only(ephemeral);
        }
        let handshake = builder.build_initiator().map_err(display_error)?;
        Ok(Self::handshake(handshake))
    }

    pub fn kk_responder(
        local_private_key: &[u8],
        local_public_key: &[u8],
        remote_public_key: &[u8],
    ) -> CoreResult<Self> {
        Self::kk_responder_inner(local_private_key, local_public_key, remote_public_key, None)
    }

    #[doc(hidden)]
    pub fn kk_responder_with_ephemeral_for_test(
        local_private_key: &[u8],
        local_public_key: &[u8],
        remote_public_key: &[u8],
        ephemeral_private_key: &[u8],
    ) -> CoreResult<Self> {
        Self::kk_responder_inner(
            local_private_key,
            local_public_key,
            remote_public_key,
            Some(ephemeral_private_key),
        )
    }

    fn kk_responder_inner(
        local_private_key: &[u8],
        local_public_key: &[u8],
        remote_public_key: &[u8],
        ephemeral_private_key: Option<&[u8]>,
    ) -> CoreResult<Self> {
        require_matching_keypair(local_private_key, local_public_key)?;
        require_key(remote_public_key, "initiator public key")?;
        let private = Zeroizing::new(local_private_key.to_vec());
        let params: NoiseParams = KK.parse().map_err(display_error)?;
        let mut builder = core_builder(params)
            .local_private_key(private.as_slice())
            .map_err(display_error)?
            .remote_public_key(remote_public_key)
            .map_err(display_error)?;
        if let Some(ephemeral) = ephemeral_private_key {
            require_key(ephemeral, "responder ephemeral private key")?;
            builder = builder.fixed_ephemeral_key_for_testing_only(ephemeral);
        }
        let handshake = builder.build_responder().map_err(display_error)?;
        Ok(Self::handshake(handshake))
    }

    fn handshake(handshake: HandshakeState) -> Self {
        Self {
            state: SessionState::Handshake(Some(handshake)),
            handshake_hash: None,
        }
    }

    pub fn write_handshake(&mut self, payload: &[u8]) -> CoreResult<Vec<u8>> {
        let handshake = self.handshake_mut()?;
        if !handshake.is_my_turn() {
            return Err("Noise handshake write attempted out of turn".into());
        }
        let mut output = vec![0_u8; MAX_NOISE_MESSAGE];
        let size = handshake
            .write_message(payload, &mut output)
            .map_err(display_error)?;
        output.truncate(size);
        self.finish_if_ready()?;
        Ok(output)
    }

    pub fn read_handshake(&mut self, message: &[u8]) -> CoreResult<Vec<u8>> {
        if message.len() > MAX_NOISE_MESSAGE {
            return Err("Noise handshake message exceeds 65535 bytes".into());
        }
        let handshake = self.handshake_mut()?;
        if handshake.is_my_turn() {
            return Err("Noise handshake read attempted out of turn".into());
        }
        let mut payload = vec![0_u8; MAX_NOISE_MESSAGE];
        let size = handshake
            .read_message(message, &mut payload)
            .map_err(display_error)?;
        payload.truncate(size);
        self.finish_if_ready()?;
        Ok(payload)
    }

    pub fn encrypt(&mut self, plaintext: &[u8]) -> CoreResult<Vec<u8>> {
        if plaintext.len() > MAX_NOISE_MESSAGE - TAG_LEN {
            return Err("Noise plaintext exceeds transport limit".into());
        }
        let transport = self.transport_mut()?;
        let mut output = vec![0_u8; plaintext.len() + TAG_LEN];
        let size = transport
            .write_message(plaintext, &mut output)
            .map_err(display_error)?;
        output.truncate(size);
        Ok(output)
    }

    pub fn decrypt(&mut self, ciphertext: &[u8]) -> CoreResult<Vec<u8>> {
        if ciphertext.len() > MAX_NOISE_MESSAGE {
            return Err("Noise ciphertext exceeds transport limit".into());
        }
        let transport = self.transport_mut()?;
        let mut output = vec![0_u8; ciphertext.len()];
        let size = transport
            .read_message(ciphertext, &mut output)
            .map_err(display_error)?;
        output.truncate(size);
        Ok(output)
    }

    pub fn handshake_hash(&self) -> CoreResult<Vec<u8>> {
        self.handshake_hash
            .map(|hash| hash.to_vec())
            .ok_or_else(|| "Noise handshake is not complete".into())
    }

    pub fn rekey_outgoing(&mut self) -> CoreResult<()> {
        self.transport_mut()?.rekey_outgoing();
        Ok(())
    }

    pub fn rekey_incoming(&mut self) -> CoreResult<()> {
        self.transport_mut()?.rekey_incoming();
        Ok(())
    }

    fn handshake_mut(&mut self) -> CoreResult<&mut HandshakeState> {
        match &mut self.state {
            SessionState::Handshake(Some(handshake)) => Ok(handshake),
            SessionState::Handshake(None) => Err("Noise handshake is transitioning".into()),
            SessionState::Transport(_) => Err("Noise handshake is already complete".into()),
            SessionState::Closed => Err("Noise session is closed".into()),
        }
    }

    fn transport_mut(&mut self) -> CoreResult<&mut TransportState> {
        match &mut self.state {
            SessionState::Transport(transport) => Ok(transport),
            SessionState::Handshake(_) => Err("Noise handshake is not complete".into()),
            SessionState::Closed => Err("Noise session is closed".into()),
        }
    }

    fn finish_if_ready(&mut self) -> CoreResult<()> {
        let ready = matches!(
            &self.state,
            SessionState::Handshake(Some(handshake)) if handshake.is_handshake_finished()
        );
        if !ready {
            return Ok(());
        }

        let handshake = match std::mem::replace(&mut self.state, SessionState::Closed) {
            SessionState::Handshake(Some(handshake)) => handshake,
            _ => return Err("Noise session transition failed".into()),
        };
        let hash: [u8; HASH_LEN] = handshake
            .get_handshake_hash()
            .try_into()
            .map_err(|_| "Noise handshake hash must be 32 bytes".to_string())?;
        let transport = handshake.into_transport_mode().map_err(display_error)?;
        self.handshake_hash = Some(hash);
        self.state = SessionState::Transport(transport);
        Ok(())
    }
}

impl Drop for CoreNoiseSession {
    fn drop(&mut self) {
        if let Some(hash) = &mut self.handshake_hash {
            hash.zeroize();
        }
        self.state = SessionState::Closed;
    }
}

fn require_key(key: &[u8], label: &str) -> CoreResult<()> {
    if key.len() != 32 {
        return Err(format!("{label} must be 32 bytes"));
    }
    Ok(())
}

fn require_matching_keypair(private_key: &[u8], public_key: &[u8]) -> CoreResult<()> {
    require_key(private_key, "local private key")?;
    require_key(public_key, "local public key")?;
    let private: [u8; 32] = private_key.try_into().expect("length checked");
    let derived = PublicKey::from(&StaticSecret::from(private));
    if derived.as_bytes().ct_eq(public_key).unwrap_u8() != 1 {
        return Err("local public key does not match local private key".into());
    }
    Ok(())
}

fn display_error(error: impl std::fmt::Display) -> String {
    error.to_string()
}

fn core_builder(params: NoiseParams) -> Builder<'static> {
    Builder::with_resolver(params, Box::new(RejectZeroResolver))
}

struct RejectZeroResolver;

impl CryptoResolver for RejectZeroResolver {
    fn resolve_rng(&self) -> Option<Box<dyn Random>> {
        DefaultResolver.resolve_rng()
    }

    fn resolve_dh(&self, choice: &DHChoice) -> Option<Box<dyn Dh>> {
        DefaultResolver
            .resolve_dh(choice)
            .map(|inner| Box::new(RejectZeroDh { inner }) as Box<dyn Dh>)
    }

    fn resolve_hash(&self, choice: &HashChoice) -> Option<Box<dyn Hash>> {
        DefaultResolver.resolve_hash(choice)
    }

    fn resolve_cipher(&self, choice: &CipherChoice) -> Option<Box<dyn Cipher>> {
        DefaultResolver.resolve_cipher(choice)
    }
}

struct RejectZeroDh {
    inner: Box<dyn Dh>,
}

impl Dh for RejectZeroDh {
    fn name(&self) -> &'static str {
        self.inner.name()
    }

    fn pub_len(&self) -> usize {
        self.inner.pub_len()
    }

    fn priv_len(&self) -> usize {
        self.inner.priv_len()
    }

    fn set(&mut self, private_key: &[u8]) {
        self.inner.set(private_key);
    }

    fn generate(&mut self, rng: &mut dyn Random) -> Result<(), Error> {
        self.inner.generate(rng)
    }

    fn pubkey(&self) -> &[u8] {
        self.inner.pubkey()
    }

    fn privkey(&self) -> &[u8] {
        self.inner.privkey()
    }

    fn dh(&self, public_key: &[u8], output: &mut [u8]) -> Result<(), Error> {
        self.inner.dh(public_key, output)?;
        let used = self.inner.dh_len();
        if output[..used].ct_eq(&vec![0_u8; used]).unwrap_u8() == 1 {
            output[..used].zeroize();
            return Err(Error::Dh);
        }
        Ok(())
    }

    fn dh_len(&self) -> usize {
        self.inner.dh_len()
    }
}

#[cfg(target_arch = "wasm32")]
mod wasm {
    use super::CoreNoiseSession;
    use wasm_bindgen::prelude::*;

    fn js_error(message: String) -> JsError {
        JsError::new(&message)
    }

    #[wasm_bindgen]
    pub struct WasmNoiseSession {
        inner: CoreNoiseSession,
    }

    #[wasm_bindgen]
    impl WasmNoiseSession {
        #[wasm_bindgen(js_name = nkInitiator)]
        pub fn nk_initiator(remote_public_key: &[u8]) -> Result<WasmNoiseSession, JsError> {
            Ok(Self {
                inner: CoreNoiseSession::nk_initiator(remote_public_key).map_err(js_error)?,
            })
        }

        #[wasm_bindgen(js_name = nkResponder)]
        pub fn nk_responder(
            local_private_key: &[u8],
            local_public_key: &[u8],
        ) -> Result<WasmNoiseSession, JsError> {
            Ok(Self {
                inner: CoreNoiseSession::nk_responder(local_private_key, local_public_key)
                    .map_err(js_error)?,
            })
        }

        #[wasm_bindgen(js_name = nkInitiatorWithEphemeralForTest)]
        pub fn nk_initiator_with_ephemeral_for_test(
            remote_public_key: &[u8],
            ephemeral_private_key: &[u8],
        ) -> Result<WasmNoiseSession, JsError> {
            Ok(Self {
                inner: CoreNoiseSession::nk_initiator_with_ephemeral_for_test(
                    remote_public_key,
                    ephemeral_private_key,
                )
                .map_err(js_error)?,
            })
        }

        #[wasm_bindgen(js_name = nkResponderWithEphemeralForTest)]
        pub fn nk_responder_with_ephemeral_for_test(
            local_private_key: &[u8],
            local_public_key: &[u8],
            ephemeral_private_key: &[u8],
        ) -> Result<WasmNoiseSession, JsError> {
            Ok(Self {
                inner: CoreNoiseSession::nk_responder_with_ephemeral_for_test(
                    local_private_key,
                    local_public_key,
                    ephemeral_private_key,
                )
                .map_err(js_error)?,
            })
        }

        #[wasm_bindgen(js_name = kkInitiator)]
        pub fn kk_initiator(
            local_private_key: &[u8],
            local_public_key: &[u8],
            remote_public_key: &[u8],
        ) -> Result<WasmNoiseSession, JsError> {
            Ok(Self {
                inner: CoreNoiseSession::kk_initiator(
                    local_private_key,
                    local_public_key,
                    remote_public_key,
                )
                .map_err(js_error)?,
            })
        }

        #[wasm_bindgen(js_name = kkResponder)]
        pub fn kk_responder(
            local_private_key: &[u8],
            local_public_key: &[u8],
            remote_public_key: &[u8],
        ) -> Result<WasmNoiseSession, JsError> {
            Ok(Self {
                inner: CoreNoiseSession::kk_responder(
                    local_private_key,
                    local_public_key,
                    remote_public_key,
                )
                .map_err(js_error)?,
            })
        }

        #[wasm_bindgen(js_name = kkInitiatorWithEphemeralForTest)]
        pub fn kk_initiator_with_ephemeral_for_test(
            local_private_key: &[u8],
            local_public_key: &[u8],
            remote_public_key: &[u8],
            ephemeral_private_key: &[u8],
        ) -> Result<WasmNoiseSession, JsError> {
            Ok(Self {
                inner: CoreNoiseSession::kk_initiator_with_ephemeral_for_test(
                    local_private_key,
                    local_public_key,
                    remote_public_key,
                    ephemeral_private_key,
                )
                .map_err(js_error)?,
            })
        }

        #[wasm_bindgen(js_name = kkResponderWithEphemeralForTest)]
        pub fn kk_responder_with_ephemeral_for_test(
            local_private_key: &[u8],
            local_public_key: &[u8],
            remote_public_key: &[u8],
            ephemeral_private_key: &[u8],
        ) -> Result<WasmNoiseSession, JsError> {
            Ok(Self {
                inner: CoreNoiseSession::kk_responder_with_ephemeral_for_test(
                    local_private_key,
                    local_public_key,
                    remote_public_key,
                    ephemeral_private_key,
                )
                .map_err(js_error)?,
            })
        }

        #[wasm_bindgen(js_name = writeHandshake)]
        pub fn write_handshake(&mut self, payload: &[u8]) -> Result<Vec<u8>, JsError> {
            self.inner.write_handshake(payload).map_err(js_error)
        }

        #[wasm_bindgen(js_name = readHandshake)]
        pub fn read_handshake(&mut self, message: &[u8]) -> Result<Vec<u8>, JsError> {
            self.inner.read_handshake(message).map_err(js_error)
        }

        pub fn encrypt(&mut self, plaintext: &[u8]) -> Result<Vec<u8>, JsError> {
            self.inner.encrypt(plaintext).map_err(js_error)
        }

        pub fn decrypt(&mut self, ciphertext: &[u8]) -> Result<Vec<u8>, JsError> {
            self.inner.decrypt(ciphertext).map_err(js_error)
        }

        #[wasm_bindgen(getter, js_name = handshakeHash)]
        pub fn handshake_hash(&self) -> Result<Vec<u8>, JsError> {
            self.inner.handshake_hash().map_err(js_error)
        }

        #[wasm_bindgen(js_name = rekeyOutgoing)]
        pub fn rekey_outgoing(&mut self) -> Result<(), JsError> {
            self.inner.rekey_outgoing().map_err(js_error)
        }

        #[wasm_bindgen(js_name = rekeyIncoming)]
        pub fn rekey_incoming(&mut self) -> Result<(), JsError> {
            self.inner.rekey_incoming().map_err(js_error)
        }
    }
}

#[cfg(not(target_arch = "wasm32"))]
mod jni_bridge;

#[cfg(test)]
mod tests {
    use super::CoreNoiseSession;
    use x25519_dalek::{PublicKey, StaticSecret};

    fn keypair(seed: u8) -> ([u8; 32], [u8; 32]) {
        let private = [seed; 32];
        let public = *PublicKey::from(&StaticSecret::from(private)).as_bytes();
        (private, public)
    }

    fn complete(initiator: &mut CoreNoiseSession, responder: &mut CoreNoiseSession) {
        let message1 = initiator.write_handshake(&[]).unwrap();
        assert!(responder.read_handshake(&message1).unwrap().is_empty());
        let message2 = responder.write_handshake(&[]).unwrap();
        assert!(initiator.read_handshake(&message2).unwrap().is_empty());
        assert_eq!(
            initiator.handshake_hash().unwrap(),
            responder.handshake_hash().unwrap()
        );
    }

    #[test]
    fn nk_round_trip_uses_one_state_machine() {
        let (responder_private, responder_public) = keypair(7);
        let mut initiator = CoreNoiseSession::nk_initiator(&responder_public).unwrap();
        let mut responder =
            CoreNoiseSession::nk_responder(&responder_private, &responder_public).unwrap();
        complete(&mut initiator, &mut responder);

        let ciphertext = initiator.encrypt(b"peach").unwrap();
        assert_eq!(responder.decrypt(&ciphertext).unwrap(), b"peach");
        let response = responder.encrypt(b"sync").unwrap();
        assert_eq!(initiator.decrypt(&response).unwrap(), b"sync");
    }

    #[test]
    fn kk_round_trip_authenticates_both_static_keys() {
        let (initiator_private, initiator_public) = keypair(11);
        let (responder_private, responder_public) = keypair(19);
        let mut initiator = CoreNoiseSession::kk_initiator(
            &initiator_private,
            &initiator_public,
            &responder_public,
        )
        .unwrap();
        let mut responder = CoreNoiseSession::kk_responder(
            &responder_private,
            &responder_public,
            &initiator_public,
        )
        .unwrap();
        complete(&mut initiator, &mut responder);
    }

    #[test]
    fn mismatched_local_keypair_is_rejected() {
        let (private, _) = keypair(2);
        let (_, wrong_public) = keypair(3);
        let result = CoreNoiseSession::nk_responder(&private, &wrong_public);
        assert!(matches!(result, Err(error) if error.contains("does not match")));
    }

    #[test]
    fn tampered_transport_frame_is_rejected() {
        let (responder_private, responder_public) = keypair(23);
        let mut initiator = CoreNoiseSession::nk_initiator(&responder_public).unwrap();
        let mut responder =
            CoreNoiseSession::nk_responder(&responder_private, &responder_public).unwrap();
        complete(&mut initiator, &mut responder);
        let mut ciphertext = initiator.encrypt(b"vault").unwrap();
        ciphertext[0] ^= 0x80;
        assert!(responder.decrypt(&ciphertext).is_err());
    }
}