/*
 *
 *    Copyright (c) 2025 Project CHIP Authors
 *
 *    Licensed under the Apache License, Version 2.0 (the "License");
 *    you may not use this file except in compliance with the License.
 *    You may obtain a copy of the License at
 *
 *        http://www.apache.org/licenses/LICENSE-2.0
 *
 *    Unless required by applicable law or agreed to in writing, software
 *    distributed under the License is distributed on an "AS IS" BASIS,
 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *    See the License for the specific language governing permissions and
 *    limitations under the License.
 */

/**
 *    @file
 *      Trusty OS based implementation of CHIP crypto primitives
 */

#include "CHIPCryptoPAL.h"
#include <lib/core/CHIPSafeCasts.h>
#include <trusty_matter.h>

namespace chip {
namespace Crypto {

static inline void from_Handle(uint64_t * p256_handler, P256KeypairContext * context)
{
    /* For Trusty OS, p256_handle used to identify the algorithm context */
    static_assert(sizeof(context) >= sizeof(uint64_t));
    memcpy(context, p256_handler, sizeof(uint64_t));
}

static inline uint64_t to_Handle(P256KeypairContext * context)
{
    uint64_t ret_handle = 0;

    memcpy(&ret_handle, context, sizeof(uint64_t));

    return ret_handle;
}

static matter::TrustyMatter & GetTrustyMatter()
{
    static matter::TrustyMatter instance;
    return instance;
}

CHIP_ERROR P256Keypair::ECDSA_sign_msg(const uint8_t * msg, const size_t msg_length, P256ECDSASignature & out_signature) const
{
    CHIP_ERROR error = CHIP_NO_ERROR;
    size_t sig_size  = 0;
    int rc           = 0;

    VerifyOrReturnError((msg != nullptr) && (msg_length > 0), CHIP_ERROR_INVALID_ARGUMENT);

    uint8_t digest[kSHA256_Hash_Length];
    uint8_t sig[kP256_ECDSA_Signature_Length_Raw];
    memset(&digest[0], 0, sizeof(digest));
    memset(&sig[0], 0, sizeof(sig));

    ReturnErrorOnFailure(Hash_SHA256(msg, msg_length, &digest[0]));

    static_assert(P256ECDSASignature::Capacity() >= kP256_ECDSA_Signature_Length_Raw, "P256ECDSASignature must be large enough");
    VerifyOrExit(mInitialized, error = CHIP_ERROR_UNINITIALIZED);

    rc = GetTrustyMatter().P256KeypairECSignMsg(to_Handle(&mKeypair), digest, kSHA256_Hash_Length, sig, sig_size);
    VerifyOrExit(rc == MATTER_ERROR_OK, error = CHIP_ERROR_INTERNAL);
    VerifyOrExit(sig_size == kP256_ECDSA_Signature_Length_Raw, error = CHIP_ERROR_INTERNAL);

    VerifyOrExit(out_signature.SetLength(kP256_ECDSA_Signature_Length_Raw) == CHIP_NO_ERROR, error = CHIP_ERROR_INTERNAL);
    memcpy(out_signature.Bytes() + 0u, sig, sig_size);

exit:
    return error;
}

CHIP_ERROR P256Keypair::ECDH_derive_secret(const P256PublicKey & remote_public_key, P256ECDHDerivedSecret & out_secret) const
{
    CHIP_ERROR error = CHIP_NO_ERROR;
    int result;
    size_t out_buf_length = 0;

    result = GetTrustyMatter().P256KeypairECDH_derive_secret(to_Handle(&mKeypair), Uint8::to_const_uchar(remote_public_key),
                                                             remote_public_key.Length(), out_secret.Bytes(), out_buf_length);
    VerifyOrExit(result == MATTER_ERROR_OK, error = CHIP_ERROR_INTERNAL);
    VerifyOrExit((out_buf_length > 0), error = CHIP_ERROR_INTERNAL);
    SuccessOrExit(error = out_secret.SetLength(out_buf_length));

exit:
    return error;
}

CHIP_ERROR P256Keypair::Initialize(ECPKeyTarget key_target)
{
    CHIP_ERROR error = CHIP_NO_ERROR;
    uint8_t public_key[kP256_PublicKey_Length];
    uint64_t p256_handle = 0;
    int rc               = 0;

    if (!memcmp(&mKeypair, kTrustyMagicNumberFabricIndex, sizeof(kTrustyMagicNumberFabricIndex)))
    {
        /* Fixed p256_handle consists of magic number and a fabric index. */
        p256_handle = to_Handle(&mKeypair);
    }
    rc = GetTrustyMatter().P256KeypairInitialize(p256_handle, public_key);
    VerifyOrExit(rc == MATTER_ERROR_OK, error = CHIP_ERROR_INTERNAL);

    from_Handle(&p256_handle, &mKeypair);
    memcpy(Uint8::to_uchar(mPublicKey), public_key, kP256_PublicKey_Length);

    mInitialized = true;

exit:
    return error;
}

CHIP_ERROR P256Keypair::Serialize(P256SerializedKeypair & output) const
{
    CHIP_ERROR error = CHIP_NO_ERROR;
    int rc           = 0;

    size_t len = output.Length() == 0 ? output.Capacity() : output.Length();
    Encoding::BufferWriter bbuf(output.Bytes(), len);
    bbuf.Put(mPublicKey, mPublicKey.Length());

    if (!memcmp(&mKeypair, kTrustyMagicNumberFabricIndex, sizeof(kTrustyMagicNumberFabricIndex)))
    {
        /* To keep confidential, private key not provided but p256_handle */
        bbuf.Put(&mKeypair, sizeof(uint64_t));
    }
    else
    {
        uint8_t privkey[kP256_PrivateKey_Length];
        uint64_t p256_handle = to_Handle(&mKeypair);
        rc                   = GetTrustyMatter().P256KeypairSerialize(p256_handle, privkey);
        VerifyOrExit(rc == MATTER_ERROR_OK, error = CHIP_ERROR_INTERNAL);
        bbuf.Put(privkey, sizeof(privkey));
        ClearSecretData(privkey, sizeof(privkey));
    }
    VerifyOrExit(bbuf.Fit(), error = CHIP_ERROR_NO_MEMORY);
    output.SetLength(bbuf.Needed());

exit:
    return error;
}

CHIP_ERROR P256Keypair::Deserialize(P256SerializedKeypair & input)
{
    CHIP_ERROR error = CHIP_NO_ERROR;
    Encoding::BufferWriter bbuf(mPublicKey, mPublicKey.Length());
    int rc               = 0;
    uint64_t p256_handle = 0;

    uint8_t * pubkey = input.Bytes();
    VerifyOrExit(input.Length() >= mPublicKey.Length(), error = CHIP_ERROR_INVALID_ARGUMENT);
    bbuf.Put(input.ConstBytes(), mPublicKey.Length());
    VerifyOrExit(bbuf.Fit(), error = CHIP_ERROR_NO_MEMORY);

    if (!memcmp(input.Bytes() + kP256_PublicKey_Length, kTrustyMagicNumberFabricIndex, sizeof(kTrustyMagicNumberFabricIndex)))
    {
        /* If the fixed p256_handle was passed in following the public key, then the private key should be interpreted as the
         * p256_handle */
        VerifyOrExit(input.Length() == mPublicKey.Length() + sizeof(uint64_t), error = CHIP_ERROR_INVALID_ARGUMENT);
        memcpy(&p256_handle, input.Bytes() + kP256_PublicKey_Length, sizeof(uint64_t));
    }
    else
    {
        uint8_t * privkey;
        VerifyOrExit(input.Length() == mPublicKey.Length() + kP256_PrivateKey_Length, error = CHIP_ERROR_INVALID_ARGUMENT);
        privkey = input.Bytes() + kP256_PublicKey_Length;
        rc = GetTrustyMatter().P256KeypairDeserialize(p256_handle, pubkey, mPublicKey.Length(), privkey, kP256_PrivateKey_Length);
        VerifyOrExit(rc == MATTER_ERROR_OK, error = CHIP_ERROR_INTERNAL);
    }

    from_Handle(&p256_handle, &mKeypair);

    mInitialized = true;

exit:
    return error;
}

void P256Keypair::Clear()
{
    if (mInitialized)
    {
        uint64_t p256_handle = to_Handle(&mKeypair);
        GetTrustyMatter().P256KeypairDestroy(p256_handle);
        mInitialized = false;
    }
}

P256Keypair::~P256Keypair()
{
    Clear();
}

CHIP_ERROR P256Keypair::NewCertificateSigningRequest(uint8_t * out_csr, size_t & csr_length) const
{
    VerifyOrReturnError(mInitialized, CHIP_ERROR_UNINITIALIZED);

    if (GetTrustyMatter().P256KeypairNewCSR(to_Handle(&mKeypair), out_csr, csr_length) != MATTER_ERROR_OK)
    {
        return CHIP_ERROR_INTERNAL;
    }

    return CHIP_NO_ERROR;
}

} // namespace Crypto
} // namespace chip