from dataclasses import dataclass
import pymongo
import os
from datetime import datetime, timedelta
from smartezlogger import logger

@dataclass
class MQ():
    def __init__(self):
        self.mongo_db = None
        self._mongo_error = None
        self._canonical_db_name = os.getenv('MQ_CANONICAL_DB', 'smartez')
        self._canonical_collection_name = os.getenv('MQ_CANONICAL_COLLECTION', 'MQ')
        self._shard_db_name = os.getenv('MQ_SHARD_DB', 'smartez_mq')
        self._shard_collection_prefix = os.getenv('MQ_SHARD_COLLECTION_PREFIX', 'MQ_')
        self._secretary_checkpoint_collection = os.getenv('MQ_SECRETARY_CHECKPOINT_COLLECTION', 'MQSecretaryCheckpoint')
        self._secretary_checkpoint_key = os.getenv('MQ_SECRETARY_CHECKPOINT_KEY', 'fanout')
        self._secretary_claim_timeout_seconds = int(os.getenv('MQ_SECRETARY_CLAIM_TIMEOUT_SECONDS', '300'))
        self._secretary_claimer_id = os.getenv('MQ_SECRETARY_CLAIMER_ID', f"{os.uname().nodename}:{os.getpid()}")
        try:
            client = pymongo.MongoClient(
                os.environ["MONGODB_URL"])
            self.mongo_db = client
            self._ensure_indexes()
        except Exception as e:
            self._mongo_error = e
            logger.log_to_console('ERROR', 'MQ::__init__', f'failed to initialize MongoDB client: {e}')

    def _mq_collection(self):
        if self.mongo_db is None:
            if self._mongo_error is not None:
                raise RuntimeError(f'MongoDB client is not initialized: {self._mongo_error}')
            raise RuntimeError('MongoDB client is not initialized')
        return self.mongo_db[self._canonical_db_name][self._canonical_collection_name]

    def _checkpoint_collection(self):
        if self.mongo_db is None:
            if self._mongo_error is not None:
                raise RuntimeError(f'MongoDB client is not initialized: {self._mongo_error}')
            raise RuntimeError('MongoDB client is not initialized')
        return self.mongo_db[self._canonical_db_name][self._secretary_checkpoint_collection]

    def _sanitize_topic(self, topic):
        value = str(topic or '').strip().lower()
        if value == '':
            value = 'default'
        chars = []
        for char in value:
            if char.isalnum() or char in ['_', '-']:
                chars.append(char)
            else:
                chars.append('_')
        return ''.join(chars)

    def _topic_collection(self, topic):
        topic_safe = self._sanitize_topic(topic)
        collection_name = f"{self._shard_collection_prefix}{topic_safe}"
        return self.mongo_db[self._shard_db_name][collection_name]

    def _ensure_collection_indexes(self, collection):
        collection.create_index([('topic', 1), ('_id', 1)])
        collection.create_index([('topic', 1), ('activity', 1), ('_id', 1)])
        collection.create_index([('activity', 1), ('user_id', 1)])

    def _ensure_indexes(self):
        try:
            collection = self._mq_collection()
            self._ensure_collection_indexes(collection)
            collection.create_index([('mq_state', 1), ('topic', 1), ('_id', 1)])
            collection.create_index([('mq_state', 1), ('mq_claimed_at', 1), ('topic', 1), ('_id', 1)])
            self._checkpoint_collection().create_index([('_id', 1)], unique=True)
        except Exception as e:
            logger.log_to_console('WARNING', 'MQ::_ensure_indexes', 'failed to ensure MQ indexes: {}'.format(e))

    def _query_with_activity(self, topic, activity=None):
        query = {'topic': topic}
        if activity is not None:
            query['activity'] = activity
        return query

    def _find_one_from_shard(self, topic, query):
        shard_collection = self._topic_collection(topic)
        try:
            self._ensure_collection_indexes(shard_collection)
        except Exception:
            pass

        return shard_collection.find_one(query, sort=[('_id', pymongo.ASCENDING)])

    def _find_many_from_shard(self, topic, query, limit=None):
        shard_collection = self._topic_collection(topic)
        try:
            self._ensure_collection_indexes(shard_collection)
        except Exception:
            pass

        cursor = shard_collection.find(query).sort('_id', pymongo.ASCENDING)
        parsed_limit = None
        if limit is not None:
            try:
                parsed_limit = int(limit)
            except Exception:
                parsed_limit = None
            if parsed_limit is not None and parsed_limit > 0:
                cursor = cursor.limit(parsed_limit)
        messages = list(cursor)
        return messages

    def producer(self, topic, activity, user_id, params, unique = False):
        ''' Interface for the producer to send messages to the queue
            Sequence:topic, activity '''
        # Issue: #1
        message: dict = {}
        message = {
            'topic': topic,
            'activity': activity,
            'user_id': user_id,
            'params': params
        }
        if unique:
            result = self._mq_collection().find_one({'activity':activity, 'user_id':user_id, 'params':params} )
            if result:
                return True
        # saving to mongo
        result = self._mq_collection().insert_one(
            message)
        return message

    def secretary_fanout_once(self, batch_size=500, topic=None):
        try:
            size = 500
            try:
                size = int(batch_size)
            except Exception:
                size = 500
            if size < 1:
                size = 1

            now = datetime.utcnow()
            claim_timeout = max(1, self._secretary_claim_timeout_seconds)
            stale_before = now - timedelta(seconds=claim_timeout)

            claimable_state_query = {
                '$or': [
                    {'mq_state': {'$exists': False}},
                    {'mq_state': 'queued'},
                    {
                        'mq_state': 'claimed',
                        'mq_claimed_at': {'$lt': stale_before}
                    }
                ]
            }
            if topic is not None:
                claimable_state_query['topic'] = topic

            candidates = list(
                self._mq_collection()
                .find(claimable_state_query)
                .sort('_id', pymongo.ASCENDING)
                .limit(size)
            )

            if not candidates:
                return {
                    'processed': 0,
                    'claimed': 0,
                    'moved': 0
                }

            claimed = 0
            moved = 0
            for candidate in candidates:
                message_id = candidate.get('_id')
                if message_id is None:
                    continue

                claim_query = {
                    '_id': message_id,
                    '$or': [
                        {'mq_state': {'$exists': False}},
                        {'mq_state': 'queued'},
                        {
                            'mq_state': 'claimed',
                            'mq_claimed_at': {'$lt': stale_before}
                        }
                    ]
                }
                claim_update = {
                    '$set': {
                        'mq_state': 'claimed',
                        'mq_claimer': self._secretary_claimer_id,
                        'mq_claimed_at': now
                    }
                }
                claim_result = self._mq_collection().update_one(claim_query, claim_update)
                if claim_result.modified_count != 1:
                    continue
                claimed += 1

                claimed_message = self._mq_collection().find_one({'_id': message_id})
                if claimed_message is None:
                    continue

                topic_value = claimed_message.get('topic', 'default')
                shard_collection = self._topic_collection(topic_value)
                try:
                    self._ensure_collection_indexes(shard_collection)
                except Exception:
                    pass

                shard_payload = dict(claimed_message)
                shard_payload.pop('mq_state', None)
                shard_payload.pop('mq_claimer', None)
                shard_payload.pop('mq_claimed_at', None)
                shard_collection.replace_one({'_id': message_id}, shard_payload, upsert=True)

                delete_result = self._mq_collection().delete_one({
                    '_id': message_id,
                    'mq_state': 'claimed',
                    'mq_claimer': self._secretary_claimer_id
                })
                if delete_result.deleted_count == 1:
                    moved += 1

            self._checkpoint_collection().update_one(
                {'_id': self._secretary_checkpoint_key},
                {
                    '$set': {
                        'updated_at': now,
                        'claimer': self._secretary_claimer_id,
                        'claimed': claimed,
                        'moved': moved
                    }
                },
                upsert=True
            )

            return {
                'processed': moved,
                'claimed': claimed,
                'moved': moved
            }
        except Exception as e:
            logger.log_to_console('ERROR', 'MQ::secretary_fanout_once', 'error during fanout: {}'.format(e))
            return {
                'processed': 0,
                'claimed': 0,
                'moved': 0,
                'error': str(e)
            }

    def consumer(self, topic, delete=False, all=False, activity=None, limit=None):
        ''' Interface for the consumer to get messages from the queue
         '''
        try:
            query = self._query_with_activity(topic, activity)
            if not all:
                message = self._find_one_from_shard(topic, query)
                if message and delete:
                    self.delete_mq(message)
                return message

            messages = self._find_many_from_shard(topic, query, limit=limit)
            if delete and messages:
                ids = [message['_id'] for message in messages if message and '_id' in message]
                if ids:
                    self._topic_collection(topic).delete_many({'_id': {'$in': ids}})

            return messages if messages else None
        except Exception as e:
            logger.log_to_console(
                    'ERROR', 'MQ::consumer', 'error consuming message:{}'.format(e))

    def delete_invalid_activity_messages(self, topic, batch_size=200):
        try:
            collection = self._mq_collection()
            size = 200
            try:
                size = int(batch_size)
            except Exception:
                size = 200
            if size < 1:
                size = 1

            query = {
                'topic': topic,
                'activity': {'$type': 'binData'}
            }
            docs = list(collection.find(query, {'_id': 1}).sort('_id', pymongo.ASCENDING).limit(size))
            if not docs:
                return 0
            ids = [doc['_id'] for doc in docs if '_id' in doc]
            if not ids:
                return 0
            result = collection.delete_many({'_id': {'$in': ids}})
            try:
                self._topic_collection(topic).delete_many({'_id': {'$in': ids}})
            except Exception:
                pass
            return int(result.deleted_count)
        except Exception as e:
            logger.log_to_console(
                'ERROR', 'MQ::delete_invalid_activity_messages',
                'error deleting invalid activity messages:{}'.format(e)
            )
            return 0

    def delete_mq(self, message):
        try:
            message_id = message.get('_id')
            message_topic = message.get('topic')
            if message_id is None or message_topic is None:
                return False

            self._topic_collection(message_topic).delete_one({'_id': message_id})
            logger.log_to_console(
                    'INFO', 'MQ::delete_mq', 'handled message succesfully:{}'.format(message))
            return True
        except Exception as e:
            logger.log_to_console(
                    'ERROR', 'MQ::delete_mq', 'error deleting message:{}, with error: {}'.format(message,e))


