/*
 * Copyright (c) 2022, Matthew Olsson <mattco@serenityos.org>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#pragma once

#include <AK/Span.h>
#include <LibPDF/ObjectDerivatives.h>

namespace Crypto::Cipher {
enum class Intent;
}

namespace PDF {

enum class CryptFilterMethod {
    None,
    V2, // RC4
    AESV2,
    AESV3,
};

class SecurityHandler : public RefCounted<SecurityHandler> {
public:
    static PDFErrorOr<NonnullRefPtr<SecurityHandler>> create(Document*, NonnullRefPtr<DictObject> encryption_dict);

    virtual ~SecurityHandler() = default;

    virtual bool try_provide_user_password(StringView password) = 0;
    virtual bool has_user_password() const = 0;

    virtual void encrypt(NonnullRefPtr<Object>, Reference reference) const = 0;
    virtual void decrypt(NonnullRefPtr<Object>, Reference reference) const = 0;
};

class StandardSecurityHandler : public SecurityHandler {
public:
    static PDFErrorOr<NonnullRefPtr<StandardSecurityHandler>> create(Document*, NonnullRefPtr<DictObject> encryption_dict);

    StandardSecurityHandler(Document*, size_t revision, ByteString const& o_entry, ByteString const& oe_entry, ByteString const& u_entry, ByteString const& ue_entry, ByteString const& perms, u32 flags, bool encrypt_metadata, size_t length, CryptFilterMethod method);

    ~StandardSecurityHandler() override = default;

    bool try_provide_user_password(StringView password_string) override;

    bool has_user_password() const override { return m_encryption_key.has_value(); }

protected:
    void encrypt(NonnullRefPtr<Object>, Reference reference) const override;
    void decrypt(NonnullRefPtr<Object>, Reference reference) const override;

private:
    void crypt(NonnullRefPtr<Object>, Reference reference, Crypto::Cipher::Intent) const;

    ByteBuffer compute_user_password_value_r2(ByteBuffer password_string);
    ByteBuffer compute_user_password_value_r3_to_r5(ByteBuffer password_string);

    bool authenticate_user_password_r2_to_r5(StringView password_string);
    bool authenticate_user_password_r6_and_later(StringView password_string);
    bool authenticate_owner_password_r6_and_later(StringView password_string);

    ByteBuffer compute_encryption_key_r2_to_r5(ByteBuffer password_string);
    bool compute_encryption_key_r6_and_later(ByteBuffer password_string);

    enum class HashKind {
        Owner,
        User,
    };
    ByteBuffer computing_a_hash_r6_and_later(ByteBuffer input, StringView input_password, HashKind);

    Document* m_document;
    size_t m_revision;
    Optional<ByteBuffer> m_encryption_key;
    ByteString m_o_entry;
    ByteString m_oe_entry;
    ByteString m_u_entry;
    ByteString m_ue_entry;
    ByteString m_perms_entry;
    u32 m_flags;
    bool m_encrypt_metadata;
    size_t m_length;
    CryptFilterMethod m_method;
};

class RC4 {
public:
    RC4(ReadonlyBytes key);

    void generate_bytes(ByteBuffer&);
    ByteBuffer encrypt(ReadonlyBytes bytes);

private:
    Array<size_t, 256> m_bytes;
};

}