use super::CoreNoiseSession;
use jni::{
    JNIEnv,
    objects::{JByteArray, JObject},
    sys::{jbyteArray, jlong},
};
use once_cell::sync::Lazy;
use std::{
    collections::HashMap,
    ptr,
    sync::{
        Mutex,
        atomic::{AtomicI64, Ordering},
    },
};

static NEXT_HANDLE: AtomicI64 = AtomicI64::new(1);
static SESSIONS: Lazy<Mutex<HashMap<i64, CoreNoiseSession>>> =
    Lazy::new(|| Mutex::new(HashMap::new()));

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

fn bytes(env: &JNIEnv<'_>, value: JByteArray<'_>) -> NativeResult<Vec<u8>> {
    env.convert_byte_array(value)
        .map_err(|error| error.to_string())
}

fn create(session: CoreNoiseSession) -> NativeResult<jlong> {
    let handle = NEXT_HANDLE.fetch_add(1, Ordering::Relaxed);
    SESSIONS
        .lock()
        .map_err(|_| "Noise session registry is unavailable".to_string())?
        .insert(handle, session);
    Ok(handle)
}

fn with_session<T>(
    handle: jlong,
    operation: impl FnOnce(&mut CoreNoiseSession) -> NativeResult<T>,
) -> NativeResult<T> {
    let mut sessions = SESSIONS
        .lock()
        .map_err(|_| "Noise session registry is unavailable".to_string())?;
    let session = sessions
        .get_mut(&handle)
        .ok_or_else(|| "Noise session handle is invalid or closed".to_string())?;
    operation(session)
}

fn throw(env: &mut JNIEnv<'_>, message: String) {
    let _ = env.throw_new("java/lang/IllegalStateException", message);
}

fn long_result(env: &mut JNIEnv<'_>, result: NativeResult<jlong>) -> jlong {
    match result {
        Ok(value) => value,
        Err(error) => {
            throw(env, error);
            0
        }
    }
}

fn array_result(env: &mut JNIEnv<'_>, result: NativeResult<Vec<u8>>) -> jbyteArray {
    match result {
        Ok(value) => match env.byte_array_from_slice(&value) {
            Ok(array) => array.into_raw(),
            Err(error) => {
                throw(env, error.to_string());
                ptr::null_mut()
            }
        },
        Err(error) => {
            throw(env, error);
            ptr::null_mut()
        }
    }
}

#[unsafe(no_mangle)]
pub extern "system" fn Java_com_peachpasswords_sync_RustNoiseNative_nativeCreateNkInitiator(
    mut env: JNIEnv,
    _object: JObject,
    remote_public_key: JByteArray,
) -> jlong {
    let result = bytes(&env, remote_public_key)
        .and_then(|key| CoreNoiseSession::nk_initiator(&key).and_then(create));
    long_result(&mut env, result)
}

#[unsafe(no_mangle)]
pub extern "system" fn Java_com_peachpasswords_sync_RustNoiseNative_nativeCreateNkResponder(
    mut env: JNIEnv,
    _object: JObject,
    local_private_key: JByteArray,
    local_public_key: JByteArray,
) -> jlong {
    let result = bytes(&env, local_private_key).and_then(|private| {
        bytes(&env, local_public_key)
            .and_then(|public| CoreNoiseSession::nk_responder(&private, &public).and_then(create))
    });
    long_result(&mut env, result)
}

#[unsafe(no_mangle)]
pub extern "system" fn Java_com_peachpasswords_sync_RustNoiseNative_nativeCreateNkInitiatorForTest(
    mut env: JNIEnv,
    _object: JObject,
    remote_public_key: JByteArray,
    ephemeral_private_key: JByteArray,
) -> jlong {
    let result = bytes(&env, remote_public_key).and_then(|remote| {
        bytes(&env, ephemeral_private_key).and_then(|ephemeral| {
            CoreNoiseSession::nk_initiator_with_ephemeral_for_test(&remote, &ephemeral)
                .and_then(create)
        })
    });
    long_result(&mut env, result)
}

#[unsafe(no_mangle)]
pub extern "system" fn Java_com_peachpasswords_sync_RustNoiseNative_nativeCreateNkResponderForTest(
    mut env: JNIEnv,
    _object: JObject,
    local_private_key: JByteArray,
    local_public_key: JByteArray,
    ephemeral_private_key: JByteArray,
) -> jlong {
    let result = bytes(&env, local_private_key).and_then(|private| {
        bytes(&env, local_public_key).and_then(|public| {
            bytes(&env, ephemeral_private_key).and_then(|ephemeral| {
                CoreNoiseSession::nk_responder_with_ephemeral_for_test(
                    &private, &public, &ephemeral,
                )
                .and_then(create)
            })
        })
    });
    long_result(&mut env, result)
}

#[unsafe(no_mangle)]
pub extern "system" fn Java_com_peachpasswords_sync_RustNoiseNative_nativeCreateKkInitiator(
    mut env: JNIEnv,
    _object: JObject,
    local_private_key: JByteArray,
    local_public_key: JByteArray,
    remote_public_key: JByteArray,
) -> jlong {
    let result = bytes(&env, local_private_key).and_then(|private| {
        bytes(&env, local_public_key).and_then(|public| {
            bytes(&env, remote_public_key).and_then(|remote| {
                CoreNoiseSession::kk_initiator(&private, &public, &remote).and_then(create)
            })
        })
    });
    long_result(&mut env, result)
}

#[unsafe(no_mangle)]
pub extern "system" fn Java_com_peachpasswords_sync_RustNoiseNative_nativeCreateKkResponder(
    mut env: JNIEnv,
    _object: JObject,
    local_private_key: JByteArray,
    local_public_key: JByteArray,
    remote_public_key: JByteArray,
) -> jlong {
    let result = bytes(&env, local_private_key).and_then(|private| {
        bytes(&env, local_public_key).and_then(|public| {
            bytes(&env, remote_public_key).and_then(|remote| {
                CoreNoiseSession::kk_responder(&private, &public, &remote).and_then(create)
            })
        })
    });
    long_result(&mut env, result)
}

#[unsafe(no_mangle)]
pub extern "system" fn Java_com_peachpasswords_sync_RustNoiseNative_nativeCreateKkInitiatorForTest(
    mut env: JNIEnv,
    _object: JObject,
    local_private_key: JByteArray,
    local_public_key: JByteArray,
    remote_public_key: JByteArray,
    ephemeral_private_key: JByteArray,
) -> jlong {
    let result = bytes(&env, local_private_key).and_then(|private| {
        bytes(&env, local_public_key).and_then(|public| {
            bytes(&env, remote_public_key).and_then(|remote| {
                bytes(&env, ephemeral_private_key).and_then(|ephemeral| {
                    CoreNoiseSession::kk_initiator_with_ephemeral_for_test(
                        &private, &public, &remote, &ephemeral,
                    )
                    .and_then(create)
                })
            })
        })
    });
    long_result(&mut env, result)
}

#[unsafe(no_mangle)]
pub extern "system" fn Java_com_peachpasswords_sync_RustNoiseNative_nativeCreateKkResponderForTest(
    mut env: JNIEnv,
    _object: JObject,
    local_private_key: JByteArray,
    local_public_key: JByteArray,
    remote_public_key: JByteArray,
    ephemeral_private_key: JByteArray,
) -> jlong {
    let result = bytes(&env, local_private_key).and_then(|private| {
        bytes(&env, local_public_key).and_then(|public| {
            bytes(&env, remote_public_key).and_then(|remote| {
                bytes(&env, ephemeral_private_key).and_then(|ephemeral| {
                    CoreNoiseSession::kk_responder_with_ephemeral_for_test(
                        &private, &public, &remote, &ephemeral,
                    )
                    .and_then(create)
                })
            })
        })
    });
    long_result(&mut env, result)
}

#[unsafe(no_mangle)]
pub extern "system" fn Java_com_peachpasswords_sync_RustNoiseNative_nativeWriteHandshake(
    mut env: JNIEnv,
    _object: JObject,
    handle: jlong,
    payload: JByteArray,
) -> jbyteArray {
    let result = bytes(&env, payload)
        .and_then(|value| with_session(handle, |session| session.write_handshake(&value)));
    array_result(&mut env, result)
}

#[unsafe(no_mangle)]
pub extern "system" fn Java_com_peachpasswords_sync_RustNoiseNative_nativeReadHandshake(
    mut env: JNIEnv,
    _object: JObject,
    handle: jlong,
    message: JByteArray,
) -> jbyteArray {
    let result = bytes(&env, message)
        .and_then(|value| with_session(handle, |session| session.read_handshake(&value)));
    array_result(&mut env, result)
}

#[unsafe(no_mangle)]
pub extern "system" fn Java_com_peachpasswords_sync_RustNoiseNative_nativeEncrypt(
    mut env: JNIEnv,
    _object: JObject,
    handle: jlong,
    plaintext: JByteArray,
) -> jbyteArray {
    let result = bytes(&env, plaintext)
        .and_then(|value| with_session(handle, |session| session.encrypt(&value)));
    array_result(&mut env, result)
}

#[unsafe(no_mangle)]
pub extern "system" fn Java_com_peachpasswords_sync_RustNoiseNative_nativeDecrypt(
    mut env: JNIEnv,
    _object: JObject,
    handle: jlong,
    ciphertext: JByteArray,
) -> jbyteArray {
    let result = bytes(&env, ciphertext)
        .and_then(|value| with_session(handle, |session| session.decrypt(&value)));
    array_result(&mut env, result)
}

#[unsafe(no_mangle)]
pub extern "system" fn Java_com_peachpasswords_sync_RustNoiseNative_nativeHandshakeHash(
    mut env: JNIEnv,
    _object: JObject,
    handle: jlong,
) -> jbyteArray {
    let result = with_session(handle, |session| session.handshake_hash());
    array_result(&mut env, result)
}

#[unsafe(no_mangle)]
pub extern "system" fn Java_com_peachpasswords_sync_RustNoiseNative_nativeRekeyOutgoing(
    mut env: JNIEnv,
    _object: JObject,
    handle: jlong,
) {
    if let Err(error) = with_session(handle, |session| session.rekey_outgoing()) {
        throw(&mut env, error);
    }
}

#[unsafe(no_mangle)]
pub extern "system" fn Java_com_peachpasswords_sync_RustNoiseNative_nativeRekeyIncoming(
    mut env: JNIEnv,
    _object: JObject,
    handle: jlong,
) {
    if let Err(error) = with_session(handle, |session| session.rekey_incoming()) {
        throw(&mut env, error);
    }
}

#[unsafe(no_mangle)]
pub extern "system" fn Java_com_peachpasswords_sync_RustNoiseNative_nativeClose(
    mut env: JNIEnv,
    _object: JObject,
    handle: jlong,
) {
    match SESSIONS.lock() {
        Ok(mut sessions) => {
            sessions.remove(&handle);
        }
        Err(_) => throw(&mut env, "Noise session registry is unavailable".into()),
    }
}