import logging
from timestream.dimensions import TimestreamDimensions
from timestream.record import TimestreamRecord
from timestream import DEFAULT_LOGGER
from typing import List, Tuple

DEFAULT_MAX_BATCH_SIZE = 100 # maximum set by AWS Timestream service limits
DEFAULT_MIN_COMMON_DIMENSIONS = 3


class Batch:
    def __init__(
        self,
        records: List[TimestreamRecord] = None,
        common_dimensions: TimestreamDimensions = None
    ):
        self.records = records or []
        self.common_dimensions = common_dimensions or TimestreamDimensions()
        


def batch_records(
    records: List[TimestreamRecord],
    max_batch_size: int = DEFAULT_MAX_BATCH_SIZE,
    min_common_dimensions: int = DEFAULT_MIN_COMMON_DIMENSIONS
) -> List[Batch]:
    batches, remaining = [], records
    DEFAULT_LOGGER.info(f"MAX BATCH || {max_batch_size}  MIN COMMON DIM || {min_common_dimensions}")
    while len(remaining) > 0:
        DEFAULT_LOGGER.info(f"Remaining records to batch || {len(remaining)}")
        batch, remaining = _batch_records_helper(
            remaining, max_batch_size, min_common_dimensions
        )
        batches.append(batch)
    return batches


def _batch_records_helper(
    records: List[TimestreamRecord],
    max_batch_size: int = DEFAULT_MAX_BATCH_SIZE,
    min_common_dimensions: int = DEFAULT_MIN_COMMON_DIMENSIONS
) -> Tuple[Batch, list]:
    assert max_batch_size > 0 and max_batch_size <= DEFAULT_MAX_BATCH_SIZE
    assert min_common_dimensions >= 0

    if len(records) < 2:
        return Batch(records, TimestreamDimensions()), []
    
    indexes = [0] # indexes for all the records in the batch
    common_dimensions = records[0].dimensions
    min_common_dimensions = min(min_common_dimensions, len(common_dimensions))
    
    for i in range(1, len(records)):
        if len(indexes) >= max_batch_size:
            break

        dimensions, shared = records[i].dimensions, {}
        for k in common_dimensions.get_dimension_keys():
            if dimensions.get_dimension(k) == common_dimensions.get_dimension(k):
                shared[k] = True
        # DEFAULT_LOGGER.info(f"SHARED {shared}")
        # Add the record to the batch if it has the minimum number of shared dimensions
        if len(shared) >= min_common_dimensions:
            DEFAULT_LOGGER.info(f"SHARED >= MCD {len(shared)}")
            keys_to_remove = [k for k in common_dimensions.get_dimension_keys() if k not in shared]
            for k in keys_to_remove:
                common_dimensions.remove_dimension(k)
            indexes.append(i)

    batch = Batch([records[i] for i in indexes], common_dimensions)
    remaining = [records[i] for i in range(len(records)) if i not in indexes]
    return batch, remaining
