"""Kafka executor class for the feed ingestion tasks."""
import json
from typing import Dict

import boto3
from kafka import KafkaAdminClient
from kafka import KafkaConsumer
from kafka import KafkaProducer

from feed_ingestion.conf.config import BOTO3_CONFIG


class KafkaExecutor:
    """Factory class to abstract Kafka operations."""

    def __init__(self, cluster_name: str):
        """Initialize executor.

        Args:
            cluster_name: the name of the kafka_cluster
        """
        self.cluster_name = cluster_name
        self.bootstrap_brokers = self._bootstrap_brokers()
        self.security_protocol = 'SSL'
        self.start_offsets = {}
        self.end_offsets = {}
        self.basic_config = dict(
            security_protocol=self.security_protocol,
            bootstrap_servers=self.bootstrap_brokers)

    def _bootstrap_brokers(self):
        client = boto3.client('kafka', config=BOTO3_CONFIG)
        cluster_list = client.list_clusters()['ClusterInfoList']
        for cluster in cluster_list:
            if self.cluster_name in cluster['ClusterName']:
                cluster_info = client.get_bootstrap_brokers(
                    ClusterArn=cluster['ClusterArn'])
                return cluster_info['BootstrapBrokerStringTls'].split(',')

    def admin_client(self, client_id: str, **kwargs) -> KafkaAdminClient:
        """Return admin client instance."""
        return KafkaAdminClient(
            client_id=client_id, **kwargs, **self.basic_config)

    def producer(
            self,
            client_id: str,
            **kwargs) -> KafkaProducer:
        """Return producer instance."""
        def _json_serializer(obj: Dict) -> bytes:
            return json.dumps(obj, default=str).encode('utf-8')

        return KafkaProducer(
            client_id=client_id,
            key_serializer=_json_serializer,
            value_serializer=_json_serializer,
            **kwargs,
            **self.basic_config)

    def consumer(
            self,
            client_id: str,
            auto_offset_reset: str,
            group_id: str = None,
            topics: str = None,
            enable_auto_commit: bool = True,
            **kwargs) -> KafkaConsumer:
        """Return consumer instance."""
        if not topics:
            return KafkaConsumer(
                client_id=client_id,
                group_id=group_id,
                auto_offset_reset=auto_offset_reset,
                enable_auto_commit=enable_auto_commit,
                **self.basic_config)
        return KafkaConsumer(
            topics,
            client_id=client_id,
            group_id=group_id,
            auto_offset_reset=auto_offset_reset,
            enable_auto_commit=enable_auto_commit,
            **kwargs,
            **self.basic_config)