# Copyright 2021-2023 VMware, Inc.
# SPDX-License-Identifier: Apache-2.0
from copy import deepcopy
from typing import Any
from typing import Optional

import pytest
from vdk.api.job_input import IIngester
from vdk.api.job_input import IProperties
from vdk.plugin.data_sources.auto_generated import (
    AutoGeneratedDataSource,
)
from vdk.plugin.data_sources.auto_generated import (
    AutoGeneratedDataSourceConfiguration,
)
from vdk.plugin.data_sources.ingester import DataSourceIngester


class MockJobInput(IIngester, IProperties):
    def __init__(self):
        self.ingested_data = []
        self.props = {}

    def send_object_for_ingestion(
        self,
        payload: dict,
        destination_table: Optional[str],
        method: Optional[str],
        target: Optional[str],
        collection_id: Optional[str] = None,
    ):
        self.ingested_data.append(payload)

    def send_tabular_data_for_ingestion(
        self,
        rows: iter,
        column_names: list,
        destination_table: Optional[str],
        method: Optional[str],
        target: Optional[str],
        collection_id: Optional[str] = None,
    ):
        raise NotImplemented("send_tabular_data_for_ingestion")

    def get_property(self, name: str, default_value: Any = None) -> str:
        return self.props.get(name, default_value)

    def get_all_properties(self) -> dict:
        return self.props

    def set_all_properties(self, properties: dict):
        self.props = deepcopy(properties)


@pytest.mark.parametrize("num_records, num_streams", [(5, 2), (10, 1)])
def test_data_source_ingester(num_records, num_streams):
    # Arrange
    data_source, data_source_ingester, mock_job_input = arrange(
        num_records, num_streams
    )

    # Act
    data_source_ingester.start_ingestion("id", data_source)
    data_source_ingester.terminate_and_wait_to_finish()

    # Assert
    assert num_records * num_streams == len(mock_job_input.ingested_data)
    assert min(mock_job_input.ingested_data, key=lambda x: x["id"])["id"] == 1
    assert max(mock_job_input.ingested_data, key=lambda x: x["id"])["id"] == num_records


@pytest.mark.parametrize("num_records, num_streams", [(2, 2)])
def test_data_source_ingester_with_duplicate_source_raises_error(
    num_records, num_streams
):
    # Arrange
    data_source, data_source_ingester, mock_job_input = arrange(
        num_records, num_streams
    )

    # Act
    data_source_ingester.ingest_data_source("auto", data_source)
    with pytest.raises(ValueError):
        data_source_ingester.ingest_data_source(
            "auto", data_source, method="second_method"
        )


@pytest.mark.parametrize("num_records, num_streams", [(5, 2), (10, 1)])
def test_data_source_ingester_with_two_sources(num_records, num_streams):
    # Arrange
    data_source, data_source_ingester, mock_job_input = arrange(
        num_records, num_streams
    )

    # Act
    data_source_ingester.ingest_data_source("auto", data_source)
    data_source_ingester.ingest_data_source(
        "auto2", data_source, method="second_method"
    )
    data_source_ingester.terminate_and_wait_to_finish()

    # Assert
    assert num_records * num_streams * 2 == len(mock_job_input.ingested_data)


@pytest.mark.parametrize("num_records, num_streams", [(5, 2), (10, 1)])
def test_data_source_ingester_with_state(num_records, num_streams):
    # First time
    mock_job_input = MockJobInput()

    data_source_ingester = DataSourceIngester(mock_job_input)
    data_source = get_data_source(num_records, num_streams)

    data_source_ingester.ingest_data_source("auto", data_source)
    data_source_ingester.terminate_and_wait_to_finish()

    # Second time:

    data_source_ingester = DataSourceIngester(mock_job_input)
    data_source = get_data_source(num_records, num_streams)

    data_source_ingester.ingest_data_source("auto", data_source)
    data_source_ingester.terminate_and_wait_to_finish()

    # Assert
    assert num_records * num_streams * 2 == len(mock_job_input.ingested_data)
    assert min(mock_job_input.ingested_data, key=lambda x: x["id"])["id"] == 1
    assert (
        max(mock_job_input.ingested_data, key=lambda x: x["id"])["id"]
        == num_records * 2
    )


def arrange(num_records, num_streams):
    mock_job_input = MockJobInput()
    data_source_ingester = DataSourceIngester(mock_job_input)
    data_source = get_data_source(num_records, num_streams)
    return data_source, data_source_ingester, mock_job_input


def get_data_source(num_records, num_streams):
    config = AutoGeneratedDataSourceConfiguration(
        num_records=num_records, include_metadata=False, num_streams=num_streams
    )
    data_source = AutoGeneratedDataSource()
    data_source.configure(config)
    return data_source
