//! Key management — upload/fetch KeyPackages and hybrid public keys.

use quicprochat_proto::bytes::Bytes;
use quicprochat_proto::prost::Message;
use quicprochat_proto::{method_ids, qpc::v1};
use quicprochat_rpc::client::RpcClient;

use crate::error::SdkError;

/// Upload a KeyPackage for pre-key distribution.
/// Returns the SHA-256 fingerprint echoed by the server.
pub async fn upload_key_package(
    rpc: &RpcClient,
    identity_key: &[u8],
    package: &[u8],
) -> Result<Vec<u8>, SdkError> {
    let req = v1::UploadKeyPackageRequest {
        identity_key: identity_key.to_vec(),
        package: package.to_vec(),
    };
    let resp_bytes = rpc
        .call(method_ids::UPLOAD_KEY_PACKAGE, Bytes::from(req.encode_to_vec()))
        .await?;
    let resp = v1::UploadKeyPackageResponse::decode(resp_bytes)
        .map_err(|e| SdkError::Other(anyhow::anyhow!("decode UploadKeyPackageResponse: {e}")))?;
    Ok(resp.fingerprint)
}

/// Fetch a KeyPackage for a peer (consumed: single-use).
/// Returns `None` if the peer has no available key packages.
pub async fn fetch_key_package(
    rpc: &RpcClient,
    identity_key: &[u8],
) -> Result<Option<Vec<u8>>, SdkError> {
    let req = v1::FetchKeyPackageRequest {
        identity_key: identity_key.to_vec(),
    };
    let resp_bytes = rpc
        .call(method_ids::FETCH_KEY_PACKAGE, Bytes::from(req.encode_to_vec()))
        .await?;
    let resp = v1::FetchKeyPackageResponse::decode(resp_bytes)
        .map_err(|e| SdkError::Other(anyhow::anyhow!("decode FetchKeyPackageResponse: {e}")))?;
    if resp.package.is_empty() {
        Ok(None)
    } else {
        Ok(Some(resp.package))
    }
}

/// Upload hybrid public key (X25519 + ML-KEM-768).
pub async fn upload_hybrid_key(
    rpc: &RpcClient,
    identity_key: &[u8],
    hybrid_public_key: &[u8],
) -> Result<(), SdkError> {
    let req = v1::UploadHybridKeyRequest {
        identity_key: identity_key.to_vec(),
        hybrid_public_key: hybrid_public_key.to_vec(),
    };
    let resp_bytes = rpc
        .call(method_ids::UPLOAD_HYBRID_KEY, Bytes::from(req.encode_to_vec()))
        .await?;
    let _resp = v1::UploadHybridKeyResponse::decode(resp_bytes)
        .map_err(|e| SdkError::Other(anyhow::anyhow!("decode UploadHybridKeyResponse: {e}")))?;
    Ok(())
}

/// Fetch a peer's hybrid public key.
/// Returns `None` if the peer has not uploaded a hybrid key.
pub async fn fetch_hybrid_key(
    rpc: &RpcClient,
    identity_key: &[u8],
) -> Result<Option<Vec<u8>>, SdkError> {
    let req = v1::FetchHybridKeyRequest {
        identity_key: identity_key.to_vec(),
    };
    let resp_bytes = rpc
        .call(method_ids::FETCH_HYBRID_KEY, Bytes::from(req.encode_to_vec()))
        .await?;
    let resp = v1::FetchHybridKeyResponse::decode(resp_bytes)
        .map_err(|e| SdkError::Other(anyhow::anyhow!("decode FetchHybridKeyResponse: {e}")))?;
    if resp.hybrid_public_key.is_empty() {
        Ok(None)
    } else {
        Ok(Some(resp.hybrid_public_key))
    }
}

/// Batch fetch hybrid keys for multiple identities.
/// Returns one `Option<Vec<u8>>` per requested identity, in the same order.
pub async fn fetch_hybrid_keys(
    rpc: &RpcClient,
    identity_keys: &[Vec<u8>],
) -> Result<Vec<Option<Vec<u8>>>, SdkError> {
    let req = v1::FetchHybridKeysRequest {
        identity_keys: identity_keys.to_vec(),
    };
    let resp_bytes = rpc
        .call(method_ids::FETCH_HYBRID_KEYS, Bytes::from(req.encode_to_vec()))
        .await?;
    let resp = v1::FetchHybridKeysResponse::decode(resp_bytes)
        .map_err(|e| SdkError::Other(anyhow::anyhow!("decode FetchHybridKeysResponse: {e}")))?;
    let result = resp
        .keys
        .into_iter()
        .map(|k| if k.is_empty() { None } else { Some(k) })
        .collect();
    Ok(result)
}