# Copyright 2021-2023 VMware, Inc.
# SPDX-License-Identifier: Apache-2.0
from datetime import datetime
from typing import Iterator
from typing import List

from vdk.plugin.data_sources.config import config_class
from vdk.plugin.data_sources.config import config_field
from vdk.plugin.data_sources.data_source import DataSourcePayload
from vdk.plugin.data_sources.data_source import IDataSource
from vdk.plugin.data_sources.data_source import (
    IDataSourceConfiguration,
)
from vdk.plugin.data_sources.data_source import IDataSourceStream
from vdk.plugin.data_sources.factory import data_source
from vdk.plugin.data_sources.state import IDataSourceState


@config_class(
    name="auto_generated", description="Configuration for Auto generated data source"
)
class AutoGeneratedDataSourceConfiguration(IDataSourceConfiguration):
    num_records: int = config_field(
        description="Number of records to return", default=2
    )
    include_metadata: bool = config_field(
        description="If true autogenerated metadata is included in the response",
        default=False,
    )
    num_streams: int = config_field(
        description="The number of streams the data source would have", default=1
    )


class AutoGeneratedDataSourceStream(IDataSourceStream):
    """
    A single data stream in AutoGeneratedDataSource
    """

    def name(self) -> str:
        return f"stream_{self._stream_number}"

    def __init__(
        self,
        config: AutoGeneratedDataSourceConfiguration,
        stream_number: int,
        start_id: int = 0,
    ):
        self._config = config
        self._stream_number = stream_number
        self._data = self._generate_data(start_id)

    def _generate_data(self, start_id: int) -> List[DataSourcePayload]:
        generated_data = []
        for i in range(self._config.num_records):
            data = {
                "id": start_id + 1 + i,
                "name": f"Stream_{self._stream_number}_Name_{i}",
                "stream": self._stream_number,
            }
            metadata = (
                {"timestamp": datetime.now()} if self._config.include_metadata else {}
            )
            generated_data.append(
                DataSourcePayload(
                    data=data, metadata=metadata, state={"last_id": data["id"]}
                )
            )
        return generated_data

    def read(self) -> Iterator[DataSourcePayload]:
        for i in range(0, len(self._data), 1):
            yield self._data[i]


@data_source(
    name="auto-generated-data", config_class=AutoGeneratedDataSourceConfiguration
)
class AutoGeneratedDataSource(IDataSource):
    """
    Data source who is only generating some dummy data for testing purposes.
    """

    def __init__(self):
        self._config = None
        self._streams = []

    def configure(self, config: AutoGeneratedDataSourceConfiguration):
        self._config = config

    def connect(self, state: IDataSourceState):
        if not isinstance(self._config, AutoGeneratedDataSourceConfiguration):
            raise RuntimeError(
                f"config type must be {AutoGeneratedDataSourceConfiguration}"
            )
        self._streams = [
            AutoGeneratedDataSourceStream(
                self._config, i, state.read_stream(f"stream_{i}").get("last_id", 0)
            )
            for i in range(self._config.num_streams)
        ]

    def disconnect(self):
        self._streams = []

    def streams(self) -> List[IDataSourceStream]:
        return self._streams
