"""Salesforce consumer."""

import json
from datetime import datetime
from typing import List

from confluent_kafka import DeserializingConsumer
from confluent_kafka import SerializingProducer
from confluent_kafka import TopicPartition
from confluent_kafka.admin import AdminClient
from confluent_kafka.cimpl import Message
from confluent_kafka.cimpl import NewTopic
import dateutil.parser as dateParser
import requests
from simple_salesforce import format_soql
from simple_salesforce import Salesforce
from simple_salesforce import SalesforceError
from tenacity import retry
from tenacity import stop_after_attempt
from tenacity import wait_exponential

import config  # noqa
from src import constants  # noqa


logger = config.logger


class SalesforceConsumer:
    """SalesForce API client that reads messages from Salesforce and saves them into Kafka."""

    def __init__(
        self,
        sf_domain: str = config.SALESFORCE_DOMAIN,
        bootstrap_servers: str = config.BOOTSTRAP_SERVERS,
        kafka_topic: str = config.KAFKA_TOPIC,
        group_id: str = config.GROUP_ID
    ):
        """Create a new instance of SalesforceConsumer."""
        self._kafka_config = {
            'bootstrap.servers': bootstrap_servers,
            'security.protocol': 'SSL',
            'group.id': group_id
        }
        self._kafka_topic = kafka_topic
        self._offset_topic = config.KAFKA_OFFSET_TOPIC

        sf_consumer_key = config.secrets_manager_client.get_cred(
            'SALESFORCE_CONSUMER_KEY')
        sf_consumer_secret = config.secrets_manager_client.get_cred(
            'SALESFORCE_CONSUMER_SECRET')
        sf_refresh_token = config.secrets_manager_client.get_cred(
            'SALESFORCE_REFRESH_TOKEN')

        oauth_url = f'https://{sf_domain}.salesforce.com/services/oauth2/token'
        payload = '&'.join([
            f'client_id={sf_consumer_key}',
            f'client_secret={sf_consumer_secret}',
            'grant_type=refresh_token',
            f'refresh_token={sf_refresh_token}'
        ])
        headers = {'content-type': 'application/x-www-form-urlencoded'}
        response = requests.request(
            'POST', oauth_url, data=payload, headers=headers)
        credentials = response.json()

        self._sf = Salesforce(
            instance_url=credentials['instance_url'],
            session_id=credentials['access_token'],
            version=config.SALESFORCE_VERSION)

        self._producer = SerializingProducer({
            **self._kafka_config,
            'value.serializer': lambda v, _: json.dumps(v).encode(),
        })
        self._offset_producer = SerializingProducer({
            **self._kafka_config,
            'value.serializer': lambda v, _: json.dumps(v).encode(),
        })
        self._consumer = DeserializingConsumer({
            **self._kafka_config,
            'value.deserializer': lambda v, _: json.loads(v),
        })
        self._admin_client = AdminClient(self._kafka_config)
        self._received_records = 0
        self._sent_records_counter = 0

        query_file = constants.QUERY_FILE_UNQUALIFIED
        with open(query_file, 'r') as query:
            self._soql_query = query.read().replace('\n', '')

    def __enter__(self):
        """Context manager enter."""
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        """Context manager exit."""
        self._producer.flush()
        logger.info(
            f'Finished. {self._received_records} records received from SalesForce. '
            f'{self._sent_records_counter} records sent to Kafka.'
        )

    def _format_string_as_dateobj(self, date: str) -> datetime:
        """Translate kafka offset date string to datetime obj for soql."""
        try:
            return dateParser.parse(date)
        except ValueError as e:
            logger.warning('Datetime format change detected in event.gdaUnqualified.offsets.')
            raise e

    @retry(
        stop=stop_after_attempt(config.RETRY_COUNT),
        wait=wait_exponential(max=config.RETRY_MAX_WAIT),
        reraise=True,
    )
    def _load_leads(self, last_lastmodifieddate: datetime) -> List[dict]:
        """Load a list of unqualified leads from Salesforce."""
        query = format_soql(self._soql_query, last_lastmodifieddate)
        records = []
        logger.info('Loading records...')
        try:
            res = self._sf.query_all(query)
            if res['totalSize']:
                records = res['records']
                logger.info(f'Total size: {res["totalSize"]}')
                logger.info(f'Number of records: {len(records)}')
                self._received_records = len(records)
            else:
                logger.info('No records found.')
        except SalesforceError as ex:
            logger.exception('Failed to query SalesForce.. retrying.')
            raise ex
        return records

    @retry(
        stop=stop_after_attempt(config.RETRY_COUNT),
        wait=wait_exponential(max=config.RETRY_MAX_WAIT),
        reraise=True,
    )
    def _create_offset_topic(self):
        logger.info(f'Checking for existing offset topic "{self._offset_topic}"..')
        existing_offset_topic = self._admin_client.list_topics(topic=self._offset_topic)
        if not existing_offset_topic.topics:
            logger.info('Offset topic does not exist. Creating it...')
            self._admin_client.create_topics(
                [
                    NewTopic(
                        self._offset_topic,
                        num_partitions=config.OFFSET_TOPIC_PARTITIONS,
                        replication_factor=config.OFFSET_TOPIC_REPLICAS,
                    )
                ]
            )
            logger.info(f'Offset topic {self._offset_topic} created.')
        else:
            logger.info(f'Offset topic {self._offset_topic} already exists.')

    @retry(
        stop=stop_after_attempt(config.RETRY_COUNT),
        wait=wait_exponential(max=config.RETRY_MAX_WAIT),
        reraise=True,
    )
    def _read_last_lastmodifieddate(self) -> datetime:
        """Read the date of the Lead most recently modified from offset topic.

        Returns: last_lastmodifieddate (datetime)
        """
        logger.info('Loading last_lastmodifieddate saved in gdaUnqualifiedLeads.offset...')
        tp = TopicPartition(self._offset_topic, 0)
        first_offset, last_offset = self._consumer.get_watermark_offsets(tp)
        self._consumer.assign([TopicPartition(self._offset_topic, 0, last_offset - 1)])

        offsets_data = self._consumer.poll(config.OFFSETS_POLL_WAIT_TIME)
        if offsets_data:
            if offsets_data.error():
                logger.error(offsets_data.error().str())
            else:
                data = offsets_data.value()
                last_lastmodifieddate_str = data.get('last_lastmodifieddate')
                if last_lastmodifieddate_str:
                    # cast lastmodifieddate from string to datetime
                    last_lastmodifieddate = self._format_string_as_dateobj(
                        last_lastmodifieddate_str
                    )
                logger.info(
                    f"Last updated/processed lead's lastModifiedDate found: {last_lastmodifieddate}"
                )
        else:
            logger.info(
                'last_lastmodifieddate NOT FOUND in gdaUnqualifiedLeads.offset topic.'
            )
            raise RuntimeError
        return last_lastmodifieddate

    @retry(
        stop=stop_after_attempt(config.RETRY_COUNT),
        wait=wait_exponential(max=config.RETRY_MAX_WAIT),
        reraise=True,
    )
    def _save_last_lastmodifieddate(self, last_lastmodifieddate):
        """Save the date of the Lead most recently modified into gdaUnqualifiedLeads.offset.

        subsequent runs of the lambda should pick up processing leads
        where it left off, aka the most recent produced message's date from
        event.gdaUnqualifiedLeads

        Args:
            last_lastmodifieddate (str): date of lead most recently modified.
            Note: Salesforce lastModified field is returned as a string.
        """
        offset_data = {
            'last_lastmodifieddate': last_lastmodifieddate}
        self._offset_producer.produce(
            self._offset_topic,
            value=offset_data,
            key=constants.LASTMODIFIEDDATE_FIELD,
        )
        self._offset_producer.flush()
        logger.info(f'LastModifiedDate {last_lastmodifieddate} saved.')

    def _delivery_report(self, err: Exception, msg: Message):
        """Kafka message delivery report handler."""
        if err is not None:
            logger.info(f'Message delivery failed: {err}')
        else:
            self._sent_records_counter += 1
            if self._sent_records_counter == self._received_records:
                logger.info(
                    f'Kafka delivery reports for all {self._received_records} records received!'
                )

    def process_leads(self):
        """Sync Salesforce leads to Kafka."""
        self._create_offset_topic()
        last_lastmodifieddate = self._read_last_lastmodifieddate()
        # load leads using ../unqualified_leads.soql which returns
        # unqualified leads ordered by most recently updated/processed first (DESC)
        leads = self._load_leads(last_lastmodifieddate)
        logger.info(
            f'========= Processing {len(leads)} unqualified lead(s) from Salesforce. ========= '
        )
        if not leads:
            # next lambda execution will continue polling from
            # existing offset value if no records are found.
            logger.info('No leads found, offset will remain unchanged.')
            return

        # process unqualified leads to Kafka topic
        for lead in leads:
            self._producer.produce(
                self._kafka_topic,
                value=lead,
                key=lead['Id'],
                on_delivery=self._delivery_report,
            )
        logger.info(f'{len(leads)} unqualified leads records sent.')
        self._producer.flush()
        # save the most recently updated/processed unqualfied lead's lastModifiedDate
        # to the offset to use in the next SF polling window
        self._save_last_lastmodifieddate(leads[0]['LastModifiedDate'])