import logging
import asyncio
import random
from synapse.module_api import ModuleApi, errors

logger = logging.getLogger(__name__)

class AutoInviteBotModule:
    def __init__(self, config, api: ModuleApi):
        self.api = api
        self.auto_invite_user_id = config.get("auto_invite_user_id")
        if not self.auto_invite_user_id:
            raise ValueError("auto_invite_user_id must be configured")
        self.api.register_third_party_rules_callbacks(
            on_new_event=self.on_new_event
        )

    @staticmethod
    def parse_config(config):
        if "auto_invite_user_id" not in config:
            logger.error("auto_invite_user_id not provided in config")
            raise Exception("auto_invite_user_id must be provided in config")
        return config

    async def on_new_event(self, event, state_events):
        if event['type'] == "m.room.member" and event['content'].get('membership') == "invite":
            room_id = event['room_id']
            inviter = event['sender']
            invitee = event['state_key']

            try:
                await self.retry_invite(self.auto_invite_user_id, room_id, inviter)
                logger.info(f"Invited bot to room {room_id} as {invitee} was invited by {inviter}")
            except Exception as e:
                logger.error(f"Failed to invite bot to room {room_id}: {str(e)}")

    async def retry_invite(self, bot_user_id, room_id, inviter, retries=5, base_delay=1):
        for attempt in range(retries):
            try:
                state = await self.api.get_room_state(room_id)
                bot_membership = state.get(("m.room.member", bot_user_id), {}).get("content", {}).get("membership")
                if bot_membership != "join":
                    await self.api.update_room_membership(
                        sender=inviter,
                        target=bot_user_id,
                        room_id=room_id,
                        new_membership="invite"
                    )
                    return  # If invite is successful, return immediately
            except errors.SynapseError as e:
                if e.errcode == 'M_LIMIT_EXCEEDED' and attempt < retries - 1:
                    wait_time = base_delay * 2 ** attempt + random.uniform(0, 1) * base_delay
                    await asyncio.sleep(wait_time)
                else:
                    raise  # Raise the error if it's not a rate limit or we've exhausted retries

def setup(config, api):
    module_config = AutoInviteBotModule.parse_config(config)
    return AutoInviteBotModule(module_config, api)