from typing import List, Mapping, Any

from airflow.models import BaseOperator
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.providers.elasticsearch.hooks.elasticsearch import ElasticsearchPythonHook
from airflow.utils.context import Context
from more_itertools import ichunked

from common.config.app import S3_ARTIFACTS_BUCKET, CSV_DELIMITER, CSV_QUOTE_CHAR, CSV_NULL_VALUE
from common.errors.es import ESError
from common.utils.types import S3Path


__all__ = ["S3ToElasticSearchOperator"]


class S3ToElasticSearchOperator(BaseOperator):
    template_fields = ("elastic_search_index_name",)

    def __init__(
            self,
            s3_data_path: S3Path,
            elastic_search_hosts: List[str],
            elastic_search_index_name: str,
            elastic_search_index_mapping: Mapping[str, Any],
            max_page_size: int = 2000,
            request_timeout: int = 60,
            aws_conn_id=S3Hook.default_conn_name,
            **kwargs
    ):
        super().__init__(**kwargs)
        self.elastic_search_hosts = elastic_search_hosts
        self.elastic_search_index_name = elastic_search_index_name
        self.elastic_search_index_mapping = elastic_search_index_mapping
        self.max_page_size: int = max_page_size
        self.request_timeout: int = request_timeout
        self.aws_conn_id = aws_conn_id
        self.s3_data_path = s3_data_path

        self.s3 = None
        self.csv = None
        self.es = None


    def pre_execute(self, context: Context):
        from common.utils.s3 import S3
        from common.utils.csv import CSV
        from common.macros.generic import get_run_key

        conn = S3Hook(aws_conn_id=self.aws_conn_id).get_conn()
        self.s3 = S3(conn, run_id=get_run_key(context["dag_run"]), bucket_name=S3_ARTIFACTS_BUCKET)

        self.csv = CSV(
            delimiter=CSV_DELIMITER,
            quote_char=CSV_QUOTE_CHAR,
            null_value=CSV_NULL_VALUE
        )

        self.es = ElasticsearchPythonHook(hosts=self.elastic_search_hosts).get_conn

        self.s3_data_path = self.s3.get_prefixed_key(self.s3_data_path)

    def _load_batch(self, dframe: list, refresh: bool = False):
        results = self.es.bulk(body=dframe, refresh=refresh, request_timeout=self.request_timeout)

        if results.get("errors"):
            errors = []
            for result in results["items"]:
                if result["create"]["status"] != 201:
                    errors.append(result)

            raise ESError(f"Bulk indexing failed. {len(errors)} errors: {errors}")

    def _load_records(self, records) -> int:
        inserted = 0
        for chunk in ichunked(records, self.max_page_size):
            dframe = []

            for record in chunk:
                dframe.append({"create": {"_index": self.elastic_search_index_name, "_id": record["id"]}})
                dframe.append(record["value"])
                inserted += 1

            self._load_batch(dframe)

        return inserted

    def _drop_index(self):
        self.log.info(f"Dropping index '{self.elastic_search_index_name}'")
        self.es.indices.delete(index=self.elastic_search_index_name, ignore_unavailable=True)

    def _create_index(self):
        self.log.info(f"Creating index '{self.elastic_search_index_name}'")
        self.es.indices.create(index=self.elastic_search_index_name, body=self.elastic_search_index_mapping)

    def execute(self, context: Context):
        total = 0
        self._drop_index()
        self._create_index()
        for path in self.s3.get_keys(self.s3_data_path):
            records = self.csv.load(self.s3.read_object(path))
            total += self._load_records(records=records)
            self.log.info(f"{total} records upserted")