from finalsa.common.models import SqsMessage, to_message_attributes
from finalsa.sns.client.interface.client import SnsClient
from datetime import datetime, timezone
from json import dumps
from typing import Union, Callable, Dict, List, Optional
from uuid import uuid4, UUID
import boto3
import logging


class SnsClientImpl(SnsClient):

    def __init__(
        self,
        get_correlation_id: Callable[[], Union[str, UUID]] = None
    ) -> None:
        self.resource = None
        self.client = None
        self.logger = logging.getLogger("finalsa.clients")
        self.cache = {}
        self.get_correlation_id = get_correlation_id
        if get_correlation_id is None:
            self.get_correlation_id = self.__class__.default_correlation_id

    def get_resource(self):
        if self.resource is None:
            self.resource = boto3.resource('sns')
        return self.resource
    
    def get_client(self):
        if self.client is None:
            self.client = boto3.client('sns')
        return self.client

    def create_topic(self, name: str):
        try:
            topic = self.get_resource().create_topic(Name=name)
            self.logger.info("Created topic %s with ARN %s.", name, topic.arn)
        except Exception as ex:
            self.logger.error("Couldn't create topic %s.", name)
            self.logger.exception(ex)
            raise ex
        return topic

    def subscription_exists(self, topic_name: str, arn: str) -> bool:
        try:
            subscriptions = self.list_subscriptions(topic_name)
            for sub in subscriptions:
                if sub['Endpoint'] == arn:
                    return True
        except Exception as e:
            self.logger.error("Subscription doesn't exist.")
            self.logger.exception(e)
        return False

    def get_all_topics(self) -> List:
        try:
            topics_iter = self.get_resource().topics.all()
            self.logger.info("Got topics.")
            return topics_iter
        except Exception as ex:
            self.logger.error("Couldn't get topics.")
            self.logger.exception(ex)
            raise ex

    def get_or_create_topic(self, name: str):
        try:
            topic = self.get_topic(name)
        except Exception as e:
            self.logger.error(
                "Couldn't get topic %s. Creating it.", name)
            self.logger.exception(e)
            topic = self.create_topic(name)
        return topic

    def get_topic(self, topic_name: str):
        if topic_name in self.cache:
            return self.cache[topic_name]
        topic = None
        try:
            topics_iter = self.get_all_topics()
            topics_iter = filter(
                lambda t: t.arn.endswith(topic_name), topics_iter)
            topics = list(topics_iter)
            if len(topics) == 0:
                raise Exception("Topic not found.")
            self.logger.info("Got topic.")
            topic = topics[0]
            self.cache[topic_name] = topic
        except Exception as e:
            self.logger.error("Couldn't get topics.")
            self.logger.exception(e)
            raise e
        else:
            return topic

    def list_subscriptions(self, topic: str) -> List:
        try:
            topic = self.get_topic(topic)
            return self.get_client().list_subscriptions_by_topic(TopicArn=topic.arn)['Subscriptions']
        except Exception as ex:
            self.logger.error("Couldn't get subscriptions.")
            self.logger.exception(ex)
            raise ex

    def subscribe(self, topic_name: str, protocol: str, endpoint: str) -> Dict:
        topic = self.get_topic(topic_name)
        try:
            subscription = self.get_client().subscribe(
                TopicArn=topic.arn,
                Protocol=protocol, Endpoint=endpoint)
            self.logger.info(
                "Subscribed %s to topic %s with protocol %s.",
                endpoint, topic.arn, protocol)
        except Exception as e:
            self.logger.error(
                "Couldn't subscribe %s to topic %s with protocol %s.",
                endpoint, topic.arn, protocol)
            self.logger.exception(e)
            raise e
        else:
            return subscription

    def publish_message(self, topic_name: str, payload: Union[Dict, SqsMessage]) -> Dict:
        correlation_id = str(self.get_correlation_id())
        message_attrs = {
            'correlation_id':  correlation_id,
            'topic': topic_name
        }
        if isinstance(payload, SqsMessage):
            payload = payload.model_dump_json()
            return self.publish(topic_name, payload, message_attrs)
        fixed_payload = payload
        if isinstance(payload, dict):
            fixed_payload = dumps(payload)
        message = SqsMessage(
            id=str(uuid4()),
            topic=topic_name,
            payload=fixed_payload,
            correlation_id=correlation_id,
            timestamp=datetime.now(timezone.utc)
        )
        return self.publish(topic_name, message.model_dump_json(), message_attrs)

    def publish(self, topic_name: str, payload: Union[Dict, str], att_dict: Optional[Dict] = {}) -> Dict:
        topic = self.get_topic(topic_name)
        if isinstance(payload, dict):
            payload = dumps(payload)
        response = topic.publish(
            Message=payload, MessageAttributes=to_message_attributes(att_dict))
        message_id = response['MessageId']
        self.logger.info(
            "Published message with attributes %s to topic %s.", att_dict,
            topic.arn)
        return {
            'id': message_id,
            'topic': topic.arn,
            'payload': payload,
        }
