# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

import base64
from contextlib import contextmanager
from datetime import timedelta
import random
import pyarrow.fs as fs
import pyarrow as pa

import pytest

encryption_unavailable = False

try:
    import pyarrow.parquet as pq
    import pyarrow.dataset as ds
except ImportError:
    pq = None
    ds = None

try:
    from pyarrow.tests.parquet.encryption import InMemoryKmsClient
    import pyarrow.parquet.encryption as pe
except ImportError:
    encryption_unavailable = True


# Marks all of the tests in this module
pytestmark = pytest.mark.dataset


FOOTER_KEY = b"0123456789112345"
FOOTER_KEY_NAME = "footer_key"
COL_KEY = b"1234567890123450"
COL_KEY_NAME = "col_key"
KEYS = {FOOTER_KEY_NAME: FOOTER_KEY, COL_KEY_NAME: COL_KEY}
EXTRA_COL_KEY = b"2345678901234501"
EXTRA_COL_KEY_NAME = "col2_key"
COLUMNS = ["year", "n_legs", "animal"]
COLUMN_KEYS = {COL_KEY_NAME: ["n_legs", "animal"]}


def create_sample_table():
    return pa.table(
        {
            "year": [2020, 2022, 2021, 2022, 2019, 2021],
            "n_legs": [2, 2, 4, 4, 5, 100],
            "animal": [
                "Flamingo",
                "Parrot",
                "Dog",
                "Horse",
                "Brittle stars",
                "Centipede",
            ],
        }
    )


def create_encryption_config(footer_key=FOOTER_KEY_NAME, column_keys=COLUMN_KEYS):
    return pe.EncryptionConfiguration(
        footer_key=footer_key,
        plaintext_footer=False,
        column_keys=column_keys,
        encryption_algorithm="AES_GCM_V1",
        # requires timedelta or an assertion is raised
        cache_lifetime=timedelta(minutes=5.0),
        data_key_length_bits=256,
    )


def create_decryption_config():
    return pe.DecryptionConfiguration(cache_lifetime=300)


def create_kms_connection_config(keys=KEYS):
    return pe.KmsConnectionConfig(
        custom_kms_conf={
            key_name: key.decode("UTF-8")
            for key_name, key in keys.items()
        }
    )


def kms_factory(kms_connection_configuration):
    return InMemoryKmsClient(kms_connection_configuration)


@contextmanager
def cond_raises(success, error_type, match):
    if success:
        yield
    else:
        with pytest.raises(error_type, match=match):
            yield


def do_test_dataset_encryption_decryption(table, extra_column_path=None):
    # use extra column key for column extra_column_path, if given
    if extra_column_path:
        keys = dict(**KEYS, **{EXTRA_COL_KEY_NAME: EXTRA_COL_KEY})
        column_keys = dict(**COLUMN_KEYS, **{EXTRA_COL_KEY_NAME: [extra_column_path]})
        extra_column_name = extra_column_path.split(".")[0]
    else:
        keys = KEYS
        column_keys = COLUMN_KEYS
        extra_column_name = None

    # define the actual test
    def assert_decrypts(
        read_keys,
        read_columns,
        to_table_success,
        dataset_success=True,
    ):
        # use all keys for writing
        write_keys = keys
        encryption_config = create_encryption_config(FOOTER_KEY_NAME, column_keys)
        decryption_config = create_decryption_config()
        encrypt_kms_connection_config = create_kms_connection_config(write_keys)
        decrypt_kms_connection_config = create_kms_connection_config(read_keys)

        crypto_factory = pe.CryptoFactory(kms_factory)
        parquet_encryption_cfg = ds.ParquetEncryptionConfig(
            crypto_factory, encrypt_kms_connection_config, encryption_config
        )
        parquet_decryption_cfg = ds.ParquetDecryptionConfig(
            crypto_factory, decrypt_kms_connection_config, decryption_config
        )

        # create write_options with dataset encryption config
        pformat = pa.dataset.ParquetFileFormat()
        write_options = pformat.make_write_options(
            encryption_config=parquet_encryption_cfg
        )

        mockfs = fs._MockFileSystem()
        mockfs.create_dir("/")

        ds.write_dataset(
            data=table,
            base_dir="sample_dataset",
            format=pformat,
            file_options=write_options,
            filesystem=mockfs,
        )

        # read without decryption config -> errors if dataset was properly encrypted
        pformat = pa.dataset.ParquetFileFormat()
        with pytest.raises(IOError, match=r"no decryption"):
            ds.dataset("sample_dataset", format=pformat, filesystem=mockfs)

        # set decryption config for parquet fragment scan options
        pq_scan_opts = ds.ParquetFragmentScanOptions(
            decryption_config=parquet_decryption_cfg
        )
        pformat = pa.dataset.ParquetFileFormat(
            default_fragment_scan_options=pq_scan_opts
        )
        with cond_raises(dataset_success, ValueError, match="Unknown master key"):
            dataset = ds.dataset("sample_dataset", format=pformat, filesystem=mockfs)
            with cond_raises(to_table_success, ValueError, match="Unknown master key"):
                assert table.select(read_columns).equals(dataset.to_table(read_columns))

        # set decryption properties for parquet fragment scan options
        decryption_properties = crypto_factory.file_decryption_properties(
            decrypt_kms_connection_config, decryption_config)
        pq_scan_opts = ds.ParquetFragmentScanOptions(
            decryption_properties=decryption_properties
        )

        pformat = pa.dataset.ParquetFileFormat(
            default_fragment_scan_options=pq_scan_opts
        )
        with cond_raises(dataset_success, ValueError, match="Unknown master key"):
            dataset = ds.dataset("sample_dataset", format=pformat, filesystem=mockfs)
            with cond_raises(to_table_success, ValueError, match="Unknown master key"):
                assert table.select(read_columns).equals(dataset.to_table(read_columns))

    # some notable column names and keys
    all_column_names = table.column_names
    encrypted_column_names = [column_name.split(".")[0]
                              for key_name, column_names in column_keys.items()
                              for column_name in column_names]
    plaintext_column_names = [column_name
                              for column_name in all_column_names
                              if column_name not in encrypted_column_names and
                              (extra_column_path is None or
                               not extra_column_path.startswith(f"{column_name}."))]
    assert len(encrypted_column_names) > 0
    assert len(plaintext_column_names) > 0
    footer_key_only = {FOOTER_KEY_NAME: FOOTER_KEY}
    column_keys_only = {key_name: key
                        for key_name, key in keys.items()
                        if key_name != FOOTER_KEY_NAME}

    # the test scenarios

    # read with footer key only, can only read plaintext columns
    assert_decrypts(footer_key_only, plaintext_column_names, True)
    assert_decrypts(footer_key_only, encrypted_column_names, False)
    assert_decrypts(footer_key_only, all_column_names, False)

    # read with all but footer key, cannot read any columns
    assert_decrypts(column_keys_only, plaintext_column_names, False, False)
    assert_decrypts(column_keys_only, encrypted_column_names, False, False)
    assert_decrypts(column_keys_only, all_column_names, False, False)

    # with footer key and one column key, all plaintext and
    # those encrypted columns that use that key, can be read
    if len(column_keys) > 1:
        for column_key_name, column_key_column_names in column_keys.items():
            for encrypted_column_name in column_key_column_names:
                # if one nested field of a column is encrypted,
                # the entire column is considered encrypted
                encrypted_column_name = encrypted_column_name.split(".")[0]

                # decrypt with footer key and one column key
                read_keys = {key_name: key
                             for key_name, key in keys.items()
                             if key_name in [FOOTER_KEY_NAME, column_key_name]}

                # that one encrypted column can only be read
                # if it is not a column path / nested field
                plaintext_and_one_success = encrypted_column_name != extra_column_name
                plaintext_and_one = plaintext_column_names + [encrypted_column_name]

                assert_decrypts(read_keys, plaintext_column_names, True)
                assert_decrypts(read_keys, plaintext_and_one, plaintext_and_one_success)
                assert_decrypts(read_keys, encrypted_column_names, False)
                assert_decrypts(read_keys, all_column_names, False)

    # with all column keys, all columns can be read
    assert_decrypts(keys, plaintext_column_names, True)
    assert_decrypts(keys, encrypted_column_names, True)
    assert_decrypts(keys, all_column_names, True)


@pytest.mark.skipif(
    encryption_unavailable, reason="Parquet Encryption is not currently enabled"
)
def test_dataset_encryption_decryption():
    do_test_dataset_encryption_decryption(create_sample_table())


@pytest.mark.skipif(
    encryption_unavailable, reason="Parquet Encryption is not currently enabled"
)
@pytest.mark.parametrize("column_name", ["list", "list.list.element"])
def test_list_encryption_decryption(column_name):
    list_data = pa.array(
        [[1, 2, 3], [4, 5, 6], [7, 8, 9], [-1], [-2], [-3]],
        type=pa.list_(pa.int32()),
    )
    table = create_sample_table().append_column("list", list_data)

    do_test_dataset_encryption_decryption(table, column_name)


@pytest.mark.skipif(
    encryption_unavailable,
    reason="Parquet Encryption is not currently enabled"
)
@pytest.mark.parametrize(
    "column_name", ["map", "map.key_value.key", "map.key_value.value"]
)
def test_map_encryption_decryption(column_name):
    map_type = pa.map_(pa.string(), pa.int32())
    map_data = pa.array(
        [
            [("k1", 1), ("k2", 2)], [("k1", 3), ("k3", 4)], [("k2", 5), ("k3", 6)],
            [("k4", 7)], [], []
        ],
        type=map_type
    )
    table = create_sample_table().append_column("map", map_data)

    do_test_dataset_encryption_decryption(table, column_name)


@pytest.mark.skipif(
    encryption_unavailable, reason="Parquet Encryption is not currently enabled"
)
@pytest.mark.parametrize("column_name", ["struct", "struct.f1", "struct.f2"])
def test_struct_encryption_decryption(column_name):
    struct_fields = [("f1", pa.int32()), ("f2", pa.string())]
    struct_type = pa.struct(struct_fields)
    struct_data = pa.array(
        [(1, "one"), (2, "two"), (3, "three"), (4, "four"), (5, "five"), (6, "six")],
        type=struct_type
    )
    table = create_sample_table().append_column("struct", struct_data)

    do_test_dataset_encryption_decryption(table, column_name)


@pytest.mark.skipif(
    encryption_unavailable,
    reason="Parquet Encryption is not currently enabled"
)
@pytest.mark.parametrize(
    "column_name",
    [
        "col",
        "col.list.element.key_value.key",
        "col.list.element.key_value.value.f1",
        "col.list.element.key_value.value.f2"
    ]
)
def test_deep_nested_encryption_decryption(column_name):
    struct_fields = [("f1", pa.int32()), ("f2", pa.string())]
    struct_type = pa.struct(struct_fields)
    struct1 = (1, "one")
    struct2 = (2, "two")
    struct3 = (3, "three")
    struct4 = (4, "four")
    struct5 = (5, "five")
    struct6 = (6, "six")

    map_type = pa.map_(pa.int32(), struct_type)
    map1 = {1: struct1, 2: struct2}
    map2 = {3: struct3}
    map3 = {4: struct4}
    map4 = {5: struct5, 6: struct6}

    list_type = pa.list_(map_type)
    list1 = [map1, map2]
    list2 = [map3]
    list3 = [map4]
    list_data = [pa.array([list1, list2, None, list3, None, None], type=list_type)]
    table = create_sample_table().append_column("col", list_data)

    do_test_dataset_encryption_decryption(table, column_name)


@pytest.mark.skipif(
    not encryption_unavailable, reason="Parquet Encryption is currently enabled"
)
def test_write_dataset_parquet_without_encryption():
    """Test write_dataset with ParquetFileFormat and test if an exception is thrown
    if you try to set encryption_config using make_write_options"""

    # Set the encryption configuration using ParquetFileFormat
    # and make_write_options
    pformat = pa.dataset.ParquetFileFormat()

    with pytest.raises(NotImplementedError):
        _ = pformat.make_write_options(encryption_config="some value")


@pytest.mark.skipif(
    encryption_unavailable, reason="Parquet Encryption is not currently enabled"
)
def test_large_row_encryption_decryption():
    """Test encryption and decryption of a large number of rows."""

    class NoOpKmsClient(pe.KmsClient):
        def wrap_key(self, key_bytes: bytes, _: str) -> bytes:
            b = base64.b64encode(key_bytes)
            return b

        def unwrap_key(self, wrapped_key: bytes, _: str) -> bytes:
            b = base64.b64decode(wrapped_key)
            return b

    row_count = 2**15 + 1
    table = pa.Table.from_arrays(
        [pa.array(
            [random.random() for _ in range(row_count)],
            type=pa.float32()
        )], names=["foo"]
    )

    kms_config = pe.KmsConnectionConfig()
    crypto_factory = pe.CryptoFactory(lambda _: NoOpKmsClient())
    encryption_config = pe.EncryptionConfiguration(
        footer_key="UNIMPORTANT_KEY",
        column_keys={"UNIMPORTANT_KEY": ["foo"]},
        double_wrapping=True,
        plaintext_footer=False,
        data_key_length_bits=128,
    )
    pqe_config = ds.ParquetEncryptionConfig(
        crypto_factory, kms_config, encryption_config
    )
    pqd_config = ds.ParquetDecryptionConfig(
        crypto_factory, kms_config, pe.DecryptionConfiguration()
    )
    scan_options = ds.ParquetFragmentScanOptions(decryption_config=pqd_config)
    file_format = ds.ParquetFileFormat(default_fragment_scan_options=scan_options)
    write_options = file_format.make_write_options(encryption_config=pqe_config)
    file_decryption_properties = crypto_factory.file_decryption_properties(kms_config)

    mockfs = fs._MockFileSystem()
    mockfs.create_dir("/")

    path = "large-row-test-dataset"
    ds.write_dataset(table, path, format=file_format,
                     file_options=write_options, filesystem=mockfs)

    file_path = path + "/part-0.parquet"
    new_table = pq.ParquetFile(
        file_path, decryption_properties=file_decryption_properties,
        filesystem=mockfs
    ).read()
    assert table == new_table

    dataset = ds.dataset(path, format=file_format, filesystem=mockfs)
    new_table = dataset.to_table()
    assert table == new_table


@pytest.mark.skipif(
    encryption_unavailable, reason="Parquet Encryption is not currently enabled"
)
def test_dataset_encryption_with_selected_column_statistics():
    table = create_sample_table()

    encryption_config = create_encryption_config()
    decryption_config = create_decryption_config()
    kms_connection_config = create_kms_connection_config()

    crypto_factory = pe.CryptoFactory(kms_factory)
    parquet_encryption_cfg = ds.ParquetEncryptionConfig(
        crypto_factory, kms_connection_config, encryption_config
    )
    parquet_decryption_cfg = ds.ParquetDecryptionConfig(
        crypto_factory, kms_connection_config, decryption_config
    )

    # create write_options with dataset encryption config
    # and specify that statistics should be enabled for a subset of columns.
    pformat = pa.dataset.ParquetFileFormat()
    write_options = pformat.make_write_options(
        encryption_config=parquet_encryption_cfg,
        write_statistics=["year", "n_legs"]
    )

    mockfs = fs._MockFileSystem()
    mockfs.create_dir("/")

    ds.write_dataset(
        data=table,
        base_dir="sample_dataset",
        format=pformat,
        file_options=write_options,
        filesystem=mockfs,
    )

    # Open Parquet files directly and check that statistics are present
    # for the expected columns.
    pq_scan_opts = ds.ParquetFragmentScanOptions(
        decryption_config=parquet_decryption_cfg
    )
    pformat = pa.dataset.ParquetFileFormat(default_fragment_scan_options=pq_scan_opts)
    dataset = ds.dataset("sample_dataset", format=pformat, filesystem=mockfs)

    for fragment in dataset.get_fragments():
        decryption_properties = crypto_factory.file_decryption_properties(
            kms_connection_config, decryption_config, fragment.path, mockfs)
        with pq.ParquetFile(
            fragment.path,
            decryption_properties=decryption_properties,
            filesystem=mockfs,
        ) as parquet_file:
            for rg_idx in range(parquet_file.metadata.num_row_groups):
                row_group = parquet_file.metadata.row_group(rg_idx)

                assert row_group.column(0).statistics is not None
                assert row_group.column(0).statistics.min == 2019
                assert row_group.column(0).statistics.max == 2022

                assert row_group.column(1).statistics is not None
                assert row_group.column(1).statistics.min == 2
                assert row_group.column(1).statistics.max == 100

                assert row_group.column(2).statistics is None