import json
from timestream.multi_measure import TimestreamRecordMulti
from timestream.record import RejectedTimestreamRecord, TimestreamRecord
from typing import List


class TimestreamLoadResult:
    """TimestreamLoadResult is an ease of use class that wraps the results of an attempted Timestream load"""

    def __init__(self):
        self.error = None
        self.__invalid_items = []  # data items that could not be parsed into records
        self.__accepted_records = []  # records accepted by Timestream
        self.__rejected_records = []  # records rejected by Timestream and their errors

    @property
    def invalid_items(self) -> List[dict]:
        return self.__invalid_items

    @property
    def accepted_records(self) -> List[TimestreamRecord]:
        return self.__accepted_records

    @property
    def rejected_records(self) -> List[RejectedTimestreamRecord]:
        return self.__rejected_records

    def add_invalid_item(self, item: dict, error: Exception):
        assert error, "Cannot add invalid item to a Timestream load result without an error"
        self.__invalid_items.append({"item": item, "error": str(error)})

    def add_accepted_record(self, record: TimestreamRecord):
        assert (
            record and type(record) is TimestreamRecord or TimestreamRecordMulti
        ), f'Cannot add an accepted record to a Timestream load result of type "{type(record)}"'
        self.__accepted_records.append(record)

    def add_accepted_records(self, records: List[TimestreamRecord]):
        assert (
            type(records) is list
        ), f'Cannot add accepted records to a Timestream load result. Expected a list, got type "{type(records)}"'
        for r in records:
            self.add_accepted_record(r)

    def add_rejected_record(self, record: RejectedTimestreamRecord):
        assert (
            record and type(record) is RejectedTimestreamRecord
        ), f'Cannot add a rejected record to a Timestream load result of type "{type(record)}"'
        self.__rejected_records.append(record)

    def add_rejected_records(self, records: List[RejectedTimestreamRecord]):
        assert (
            type(records) is list
        ), f'Cannot add rejected records to a Timestream load result. Expected a list, got type "{type(records)}"'
        for r in records:
            self.add_rejected_record(r)

    def to_dict(self) -> dict:
        return {
            "error": str(self.error) if self.error else None,
            "invalid_items": self.__invalid_items,
            "accepted_records": [r.to_dict() for r in self.__accepted_records],
            "rejected_records": [r.to_dict() for r in self.__rejected_records],
        }

    def __repr__(self) -> str:
        return json.dumps(self.to_dict())
