from collections import defaultdict

from airflow.models import BaseOperator
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.utils.context import Context

from common.config.app import CSV_DELIMITER, CSV_QUOTE_CHAR, CSV_NULL_VALUE, S3_ARTIFACTS_BUCKET, CSV_NORMALIZERS
from common.errors.s3 import S3DataError
from common.macros.generic import get_run_key
from common.utils.csv import CSV
from common.utils.json import format_json
from common.utils.s3 import S3
from common.utils.types import S3Path


__all__ = ["TransformAggregatesOperator"]


class TransformAggregatesOperator(BaseOperator):
    ui_color = '#ededed'

    def __init__(self, batch_index: int, batch_size: int, raw_data_path: S3Path, agg_data_path: S3Path,
                 aws_conn_id: str = S3Hook.default_conn_name, **kwargs):
        super().__init__(**kwargs)
        self.batch_index = batch_index
        self.batch_size = batch_size
        self.aws_conn_id = aws_conn_id
        self.raw_data_path = raw_data_path
        self.agg_data_path = agg_data_path

    def execute(self, context: Context):
        s3_hook = S3Hook(aws_conn_id=self.aws_conn_id).get_conn()
        s3 = S3(s3_hook, run_id=get_run_key(context["dag_run"]), bucket_name=S3_ARTIFACTS_BUCKET)
        csv = CSV(
            delimiter=CSV_DELIMITER,
            quote_char=CSV_QUOTE_CHAR,
            null_value=CSV_NULL_VALUE,
            normalizers=CSV_NORMALIZERS
        )

        raw_data_path = s3.get_prefixed_key(self.raw_data_path)
        if not s3.is_data_available(raw_data_path):
            raise S3DataError(f"Data is not available at '{raw_data_path}'")

        self.log.info(f"Processing aggregates data from '{raw_data_path}'")
        mpu_key = s3.get_prefixed_key(self.agg_data_path)
        with s3.get_multipart_uploader(key=mpu_key) as mpu:
            for index, _part_index, key in s3.get_batched_keys(
                prefix=raw_data_path,
                batch_size=self.batch_size,
                batch_index=self.batch_index
            ):
                raw_data = defaultdict(list)
                for row in csv.load(s3.read_object(key=key)):
                    raw_data[row.pop("id")].append(row)

                upload_data = []
                for id_, data_row in raw_data.items():
                    upload_data.append({"id": id_, "value": format_json(data_row)})

                mpu.add_data(csv.dump(upload_data, header=("id", "value"), write_header=not index))