botbot/matrix_service/main.py

import os
import logging
from dotenv import load_dotenv

import asyncio
import httpx

from nio import AsyncClient, AsyncClientConfig, MatrixRoom, RoomMessageText, InviteMemberEvent
from nio.responses import LoginResponse


# --- Load environment variables ---
load_dotenv()
MATRIX_HOMESERVER_URL        = os.getenv("MATRIX_HOMESERVER_URL")
MATRIX_USER_ID               = os.getenv("MATRIX_USER_ID")
MATRIX_PASSWORD              = os.getenv("MATRIX_PASSWORD")
MATRIX_LOGIN_TRIES           = int(os.getenv("MATRIX_LOGIN_TRIES", 5))
MATRIX_LOGIN_DELAY_INCREMENT = int(os.getenv("MATRIX_LOGIN_DELAY_INCREMENT", 5))

LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper()

AI_HANDLER_URL   = os.getenv("AI_HANDLER_URL")
AI_HANDLER_TOKEN = os.getenv("AI_HANDLER_TOKEN")



# --- Logging Setup ---
numeric_level = getattr(logging, LOG_LEVEL, logging.INFO)
logging.basicConfig(
    level=numeric_level,
    format="%(asctime)s %(levelname)s %(name)s: %(message)s"
)
logger = logging.getLogger(__name__)


async def main() -> None:

    async def trust_all_devices(client) -> None:
        """
        Mark every other user's device as verified so we can receive their
        future Megolm keys without being blocked.
        """
        for room_id in client.rooms:
            try:
                devices_in_room = client.room_devices(room_id)
                for user_id, user_devices in devices_in_room.items():
                    if user_id == client.user_id:
                        continue
                    for dev_id, device in user_devices.items():
                        if not client.device_store.is_device_verified(device):
                            logger.info(f"Trusting {dev_id} for {user}")
                            client.verify_device(device)
            except Exception:
                logger.exception(f"Error trusting devices in {room_id}")

    async def on_invite(room, event):
        """
        Handle an invite event by joining the room and trusting all devices.
        """
        if isinstance(event, InviteMemberEvent):
            await client.join(room.room_id)
            logger.info("Joined %s", room.room_id)
            await trust_all_devices(client)

    async def send_message(room: MatrixRoom, message: str) -> None:
        """
        Send a message to `room`.
        """
        try:
            await trust_all_devices(client)

            await client.share_group_session(
                room.room_id,
                ignore_unverified_devices=True
            )

            await client.room_send(
                room_id=room.room_id,
                message_type="m.room.message",
                content={
                    "msgtype": "m.text",
                    "body": message
                },
                ignore_unverified_devices=True
            )
            logger.info("Sent message to %s: %s", room.room_id, message)
        except Exception as e:
            logger.error(f"Error sending message: {e}")


    async def on_message(room: MatrixRoom, event: RoomMessageText) -> None:
        """
        Handle incoming messages.
        """

        # Check if the message is from the bot itself
        if event.sender == client.user_id:            
            return
        
        logger.info("Received '%s' from %s in %s", event.body.strip(), event.sender, room.display_name)

        
        if isinstance(event, RoomMessageText):
            logger.info(f"Received message in {room.room_id}: {event.body}")
            payload = {
                "roomId":           event.room_id,
                "userId":           event.sender,
                "eventId":          event.event_id,
                "serverTimestamp":  event.server_timestamp,
                "content":          event.body
            }
            headers = {"Authorization": f"Bearer {AI_HANDLER_TOKEN}"}
            async with httpx.AsyncClient() as http:
                try:
                    resp = await http.post(f"{AI_HANDLER_URL}/api/v1/message", json=payload, headers=headers)
                    resp.raise_for_status()
                    data = resp.json()
                    if data.get("reply"):
                        await trust_all_devices(client)
                        await send_message(room, data["reply"])
                        logger.info("Reply sent: %s", data["reply"])
                except httpx.HTTPStatusError as e:
                    logger.error(f"HTTP error: {e.response.status_code} - {e.response.text}")
                except Exception:
                    logger.exception("Error while calling AI handler")

    # --- Initialize the client ---
    # Create the data directory if it doesn't exist
    try:
        os.makedirs("/app/data", exist_ok=True)
    except Exception as e:
        logger.error(f"Error creating data directory: {e}")
        return
    
    # Initialize the client
    config = AsyncClientConfig(store_sync_tokens=True, encryption_enabled=True)
    client = AsyncClient(MATRIX_HOMESERVER_URL, MATRIX_USER_ID, store_path="/app/data", config=config)    
  
    for i in range(MATRIX_LOGIN_TRIES):
        try:
            login_response=await client.login(password=MATRIX_PASSWORD)
            break
        except Exception as e:
            logger.error("Login failed: %s", login_response)
            logger.error(f"Login attempt {i+1} failed: {e}")
            if i == MATRIX_LOGIN_TRIES - 1:
                return
            await asyncio.sleep(MATRIX_LOGIN_DELAY_INCREMENT * (i + 1))
    logger.info("Logged in successfully")

    logger.debug("Upload one time Olm keys")
    try:
        await client.keys_upload()
    except Exception as e:
        logger.error(f"Error uploading keys: {e}")
        return
    
    await trust_all_devices(client)

    client.add_event_callback(on_invite, InviteMemberEvent)
    client.add_event_callback(on_message, RoomMessageText)

    logger.info("Starting sync loop")
    await client.sync_forever(timeout=30000) # timeout should be moved to Variable



if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("Shutting down")
        asyncio.run(client.close())