src/controllers/_bot_controller.py

import os
import re
from io import BytesIO

import aiohttp
from fastapi import APIRouter, BackgroundTasks, Form, HTTPException, Request, UploadFile
from fastapi.responses import JSONResponse, RedirectResponse
from pydantic import BaseModel
from slack_sdk.web.async_client import AsyncWebClient

from src.models import ContentDeliveryFrequency, User
from src.models._slack_team import SlackTeam
from src.services import FileService, SlackBotService, SlackTeamService, UserService
from src.utils import Helper, logger

file_set = set()
new_user_set = set()


class BotController:

    def __init__(self):
        self.router = APIRouter(prefix="/bot", tags=["BOT"])
        self.router.add_api_route("/event", self.handle_event, methods=["POST"])
        self.router.add_api_route("/auth", self.authenticate, methods=["GET"])
        self.router.add_api_route("/frequency", self.update_frequency, methods=["POST"])
        self.router.add_api_route("/install", self.install, methods=["GET"])
        self.router.add_api_route("/refresh", self.refreshContent, methods=["POST"])

    def get_welcome_message(self, firstname):
        """
        Get the welcome message for the user

        Args:
            firstname (str): The user's first name
        Returns:
            str: The welcome message
        """
        return f"""Hi {firstname} 👋\n\nWelcome to the Mandala Slackbot! Every week, you'll receive a learning tailored from the Resilient Leadership program to support you on your leadership journey. If you'd like to make this content even more personalized, share your Resilient Leadership guidebook with us either as an attachment or a link to the slide deck.\n\nOtherwise, we'll plan on sharing our curated weekly insights to keep you growing. 🙂\n\nYou can refresh the content at any time by simply typing */refresh* and sending that guidebook back our way. Looking forward to getting started! Your first learning will arrive next week.\n\nWarmly,\nTeam Mandala"""

    async def process_event_background(self, event: dict):
        """
        Process the event in the background

        Args:
            event (dict): The event to process
        Returns:
            None
        """
        file_ids = []
        global file_set
        global new_user_set
        send_message = True
        slack_user_id = None
        try:
            actual_event = event["event"]
            if "files" in actual_event:
                file_ids = [file.get("id") for file in actual_event.get("files", [])]
                if all(file_id in file_set for file_id in file_ids):
                    return
                if any(file_id in file_set for file_id in file_ids):
                    send_message = False
                file_set.update(file_ids)
            if "file" in actual_event:
                file_ids = [actual_event.get("file").get("id")]
                if any(file_id in file_set for file_id in file_ids):
                    return
                file_set.update(file_ids)

            slack_user_id = actual_event.get("user") or actual_event.get("user_id")
            async with UserService() as user_service:
                user = await user_service.get_user_by_slack_id(slack_user_id)
                if (
                    not user
                    and actual_event.get("type") == "app_home_opened"
                    and actual_event.get("tab") == "home"
                ):
                    if slack_user_id in new_user_set:
                        return
                    new_user_set.add(slack_user_id)
                    user = await self.handle_new_user(
                        event, slack_user_id, user_service
                    )
                    if not user:
                        raise Exception("Failed to register user")
                if not user:
                    logger.warning(f"User not found with slack id: {slack_user_id}")
                    return

                if "files" in actual_event or "file" in actual_event:
                    if send_message:
                        await user_service.send_message(
                            user.id,
                            "You're very welcome! 🎉 We'll be sharing these insights with you gradually over the next few weeks. 😊",
                        )
                    await self.handle_files(event, actual_event, user, user_service)
        except Exception as e:
            logger.error(f"Background task error: {e}")
            logger.exception(e)
        finally:
            file_set.difference_update(file_ids)
            new_user_set.discard(slack_user_id)

    async def handle_new_user(self, event, slack_user_id, user_service):
        """
        Register a new user

        Args:
            event (dict): The event
            slack_user_id (str): The slack user id
            user_service (UserService): The user service
        Returns:
            User: The registered user
        """
        async with SlackTeamService() as slack_team_service:
            team_id = event.get("team_id")
            logger.info(f"Registering new user {slack_user_id} in team {team_id}")
            team = await slack_team_service.get_team_by_slack_team_id(team_id)
            if not team:
                raise Exception("Team not found")

        async with SlackBotService() as slack_bot_service:
            user_email = await slack_bot_service.get_user_email(
                team.token, slack_user_id
            )
            if user_email:
                user = await user_service.get_user_info_by_email(user_email)
                if user and user.id and not user.slack_id:
                    updated_user = await user_service.update_user_info(
                        user.id, {"slack_team_id": team.id, "slack_id": slack_user_id}
                    )
                    welcome_message = self.get_welcome_message(user.name.split()[0])
                    await user_service.send_message(updated_user.id, welcome_message)
                    return updated_user
        raise Exception("Failed to register user")

    async def handle_files(self, event, actual_event, user, user_service):
        """
        Handle files uploaded by the user

        Args:
            event (dict): The event
            actual_event (dict): The actual event
            user (User): The user
            user_service (UserService): The user service
        Returns:
            None
        """
        try:
            async with SlackTeamService() as slack_team_service:
                team_id = event.get("team_id")
                team: SlackTeam = await slack_team_service.get_team_by_slack_team_id(
                    team_id
                )
                if not team:
                    raise Exception("Team not found")

            async with FileService() as file_service:
                async with SlackBotService() as slack_bot_service:
                    files = actual_event.get("files", [])
                    if not files:
                        file = (
                            actual_event.get("file") if "file" in actual_event else None
                        )
                        if not file.get("id"):
                            raise Exception("No file found")
                        file_info = await file_service.get_file_info(
                            team.token, file.get("id")
                        )
                        logger.info(f"Processing {file_info}")
                        files = [file_info.get("file")]
                    for file in files:
                        await self.process_file(
                            file,
                            team,
                            user,
                            file_service,
                            slack_bot_service,
                            user_service,
                        )
        except Exception as e:
            logger.error(f"Error handling files: {e}")
            logger.exception(e)
            await user_service.send_message(
                user.id,
                "We encountered an error processing your files. Please try again later or contact support.",
            )

    async def process_file(
        self, file, team, user, file_service, slack_bot_service, user_service
    ):
        """
        Process the file uploaded by the user

        Args:
            file (dict): The file
            team (SlackTeam): The team
            user (User): The user
            file_service (FileService): The file service
            slack_bot_service (SlackBotService): The slack bot service
            user_service (UserService): The user service
        Returns:
            None
        """
        try:
            file_id = file["id"]
            logger.info(f"Processing file id: {file_id}")

            uploaded_file = await slack_bot_service.save_file(team, file)
            if uploaded_file:
                return await file_service.process_file(
                    file=uploaded_file, user=user, slack_file_id=file_id
                )
            raise Exception(f"Failed to save file")
        except Exception as e:
            logger.error(f"Error processing file {file_id}: {e}")
            await user_service.send_message(
                user.id,
                "We encountered an error processing your files. Please try again later or contact support.",
            )

    async def handle_event(self, event: dict, background_tasks: BackgroundTasks):
        """
        Handle the event

        Args:
            event (dict): The event
            background_tasks (BackgroundTasks): The background tasks
        Returns:
            dict: The response
        """
        if "challenge" in event:
            return {"challenge": event["challenge"]}
        logger.info(f"Received event: {event}")
        if "event" not in event:
            return
        background_tasks.add_task(self.process_event_background, event)
        return self.ephemeral_response("Processing the event...")

    async def authenticate(
        self, code: str, state: str, background_tasks: BackgroundTasks
    ):
        """
        Authenticate the user

        Args:
            code (str): The code
            state (str): The state
            background_tasks (BackgroundTasks): The background tasks
        Returns:
            JSONResponse: The response
        """
        try:
            async with SlackBotService() as slack_bot_service:
                auth_response = await slack_bot_service.handle_oauth_callback(
                    code, state
                )
                if "access_token" in auth_response:
                    background_tasks.add_task(
                        self.register_slack_team,
                        slack_team_id=auth_response["team"]["id"],
                        team_name=auth_response["team"]["name"],
                        token=auth_response["access_token"],
                    )
                    return JSONResponse(content={"message": "Installed successfully"})
                else:
                    return JSONResponse(content={"message": "Failed to install"})

        except HTTPException as e:
            logger.warning(e)
            raise e
        except Exception as e:
            logger.error(e)
            raise HTTPException(status_code=500, detail=str(e))

    async def update_frequency(
        self, request: Request, background_tasks: BackgroundTasks
    ):
        """
        Update the user's content delivery frequency

        Args:
            request (Request): The request
            background_tasks (BackgroundTasks): The background tasks
        Returns:
            dict: The response
        """
        form_data = await request.form()
        payload = dict(form_data)
        background_tasks.add_task(self.update_user_frequency, payload)
        return self.ephemeral_response(
            "Your content delivery frequency is being updated"
        )

    async def install(self):
        """
        Install the app

        Returns:
            RedirectResponse: The redirect response
        """
        try:
            async with SlackBotService() as slack_bot_service:
                auth_url = await slack_bot_service.generate_auth_url()
                return RedirectResponse(url=auth_url)
        except HTTPException as e:
            logger.warning(e)
            raise e
        except Exception as e:
            logger.error(e)
            raise HTTPException(status_code=500, detail=str(e))

    async def update_user_frequency(self, payload: dict):
        """
        Update the user's content delivery frequency

        Args:
            payload (dict): The payload
        Returns:
            None
        """
        try:
            logger.info(f"Received command: {payload}")

            if payload.get("command") != "/frequency":
                return self.ephemeral_response(
                    "Invalid command. Please use /frequency command."
                )

            user_id = payload.get("user_id")
            frequency = payload.get("text", "").upper()
            new_frequency = self.get_new_frequency(frequency)
            if not new_frequency:
                valid_frequencies = ", ".join(
                    [freq.name for freq in ContentDeliveryFrequency]
                )
                return self.ephemeral_response(
                    f"Invalid frequency. Please use one of: {valid_frequencies}"
                )

            async with UserService() as user_service:
                user = await user_service.get_user_by_slack_id(user_id)
                if not user:
                    return self.ephemeral_response(
                        "User not found. Please make sure you're registered."
                    )

                updated_user = await user_service.update_user_frequency(
                    user.id, new_frequency
                )
                if not updated_user:
                    return self.ephemeral_response(
                        "Failed to update frequency. Please try again."
                    )

                next_delivery = updated_user.next_content_delivery_date.strftime(
                    "%Y-%m-%d"
                )
                await user_service.send_message(
                    user_id=user.id,
                    message=f"Your content delivery frequency has been updated to {frequency}. Next delivery scheduled for {next_delivery}",
                )

        except Exception as e:
            logger.error(f"Error updating frequency: {e}")
            logger.exception(e)
            await user_service.send_message(
                user_id=user.id,
                message="Failed to update frequency. Please try again or contact support.",
            )

    def get_new_frequency(self, frequency):
        """
        Get the new frequency

        Args:
            frequency (str): The frequency
        Returns:
            ContentDeliveryFrequency: The new frequency
        """
        return ContentDeliveryFrequency[frequency]

    async def register_slack_team(self, slack_team_id: str, team_name: str, token: str):
        """
        Register the slack team

        Args:
            slack_team_id (str): The slack team id
            team_name (str): The team name
            token (str): The token
        Returns:
            None
        """
        try:
            async with SlackTeamService() as slack_team_service:
                team = await slack_team_service.get_team_by_slack_team_id(slack_team_id)
                if not team:
                    team = await slack_team_service.register_slackTeam(
                        {
                            "name": team_name,
                            "slack_team_id": slack_team_id,
                            "token": token,
                        }
                    )
                team = await slack_team_service.update_team_info(
                    team.id, {"token": token}
                )
                logger.info(f"Registered team: {team.__dict__}")
                await self.register_team_users(team.id, token)
        except Exception as e:
            logger.error(f"Error registering team: {e}")
            logger.exception(e)

    async def register_team_users(self, slack_team_id: str, token: str):
        """
        Register the team users

        Args:
            slack_team_id (str): The slack team id
            token (str): The token
        Returns:
            None
        """
        async with SlackBotService() as slack_bot_service:
            team_users = await slack_bot_service.get_users_list(token)
            logger.info(f"got team users for team {slack_team_id}")
            for slack_user in team_users:
                await self.register_user(slack_user, slack_team_id)

    async def register_user(self, slack_user, slack_team_id):
        """
        Register the user

        Args:
            slack_user (dict): The slack user
            slack_team_id (str): The slack team id
        Returns:
            None
        """
        async with UserService() as user_service:
            user = await user_service.get_user_info_by_email(
                slack_user["profile"]["email"]
            )
            if user is not None:
                logger.info(f"got user {user.__dict__}")
                if user.id and not user.slack_id:
                    user = await user_service.update_user_info(
                        user.id,
                        {"slack_team_id": slack_team_id, "slack_id": slack_user["id"]},
                    )
                    logger.info(f"updated user {user.__dict__}")
                    logger.info(f"Sending welcome message to user {user.name}")
                    welcome_message = self.get_welcome_message(user.name.split()[0])
                    await user_service.send_message(user.id, welcome_message)

    def ephemeral_response(self, text):
        """
        Get the ephemeral response

        Args:
            text (str): The text
        Returns:
            dict: The response
        """
        return {"response_type": "ephemeral", "text": text}

    async def refreshContent(
        self,
        background_tasks: BackgroundTasks,
        command: str = Form(...),
        text: str = Form(...),
        user_id: str = Form(...),
        team_id: str = Form(...),
    ):
        """
        Refresh the user's content

        Args:
            background_tasks (BackgroundTasks): The background tasks
            command (str): The command
            text (str): The text
            user_id (str): The user id
            team_id (str): The team id
        Returns:
            dict: The response
        """
        try:
            payload = {
                "command": command,
                "text": text,
                "user_id": user_id,
                "team_id": team_id,
            }

            if command != "/refresh":
                return {
                    "response_type": "ephemeral",
                    "text": "Invalid command. Please use /refresh command.",
                }

            if not text:
                return {
                    "response_type": "ephemeral",
                    "text": "Please provide a valid file URL with the /refresh command.",
                }

            async with Helper() as helper:
                if not helper.is_valid_url(text):
                    return {
                        "response_type": "ephemeral",
                        "text": "Please provide a valid URL (e.g., https://example.com/file.pdf).",
                    }

            user_id = payload.get("user_id")
            logger.info(f"Received refresh command from user: {user_id}")
            async with UserService() as user_service:
                user = await user_service.get_user_by_slack_id(user_id)
                logger.info(f"User: {user}")
                if not user:
                    return {
                        "response_type": "ephemeral",
                        "text": "Please contact support to register your account.",
                    }

            background_tasks.add_task(
                self.refreshUserContent, payload, user, user_service
            )
            return {
                "response_type": "ephemeral",
                "text": "You're very welcome! 🎉 We'll be sharing these insights with you gradually over the next few weeks. 😊",
            }
        except Exception as e:
            logger.error(f"Error processing refresh command: {e}")
            logger.exception(e)
            return {
                "response_type": "ephemeral",
                "text": "Failed to process refresh command. Please try again or contact support.",
            }

    async def refreshUserContent(
        self, payload: dict, user: User, user_service: UserService
    ):
        """
        Refresh the user's content

        Args:
            payload (dict): The payload
            user (User): The user
            user_service (UserService): The user service
        Returns:
            None
        """
        try:
            async with SlackTeamService() as slack_team_service:
                team = await slack_team_service.get_team_by_slack_team_id(
                    payload.get("team_id")
                )
                if not team:
                    logger.error("Team not found")

                async with FileService() as file_service:
                    try:
                        # Download file from URL with timeout
                        file_url = payload.get("text")
                        async with aiohttp.ClientSession() as session:
                            async with session.get(file_url, timeout=30) as response:
                                if response.status != 200:
                                    logger.error(
                                        f"Failed to download file: {response.status}"
                                    )
                                    await user_service.send_message(
                                        user_id=user.id,
                                        message="We couldn't download the file. Please make sure the URL is accessible and try again.",
                                    )
                                    return

                                # Check content type
                                content_type = response.headers.get("content-type", "")

                                if not any(
                                    supported_type in content_type.lower()
                                    for supported_type in [
                                        "pdf",
                                        "powerpoint",
                                        "document",
                                        "application",
                                    ]
                                ):
                                    logger.error(
                                        f"Unsupported content type: {content_type}"
                                    )
                                    await user_service.send_message(
                                        user_id=user.id,
                                        message="Unsupported content type. Please provide a valid file URL.",
                                    )
                                    return

                                file_content = await response.read()
                                if not file_content:
                                    logger.error("Empty file content")
                                    await user_service.send_message(
                                        user_id=user.id,
                                        message="We couldn't download the file. Please make sure the file has content and try again.",
                                    )
                                    return

                                content_disposition = response.headers.get(
                                    "content-disposition"
                                )
                                if content_disposition:
                                    # Try to extract filename from content-disposition
                                    match = re.search(
                                        r'filename="(.+?)"', content_disposition
                                    )
                                    if match:
                                        file_name = match.group(1)
                                    else:
                                        file_name = file_url.split("/")[-1]
                                else:
                                    file_name = file_url.split("/")[-1]

                                uploaded_file = UploadFile(
                                    filename=file_name,
                                    file=BytesIO(file_content),
                                )

                        # Process the file
                        await file_service.process_file(
                            file=uploaded_file,
                            user=user,
                            slack_file_id=None,
                        )

                    except Exception as e:
                        logger.error(f"Error processing file: {e}")
                        await user_service.send_message(
                            user_id=user.id,
                            message="Failed to process file. Please try again.",
                        )

        except Exception as e:
            logger.error(f"Error processing refresh command: {e}")
            logger.exception(e)