simbarag/blueprints/imessage/__init__.py

import asyncio
import os
import hmac
import logging
import functools
import time
from collections import defaultdict

import httpx
from quart import Blueprint, request, jsonify

from blueprints.users.models import User
from blueprints.conversation.logic import (
    get_conversation_for_channel,
    add_message_to_conversation,
    persist_agent_tool_messages,
)
from blueprints.conversation.models import Conversation, ConversationMessage
from blueprints.conversation.agents import main_agent
from blueprints.conversation.memory import build_system_prompt_with_memories
from utils.strip_markdown import extract_agent_text, strip_markdown

imessage_blueprint = Blueprint("imessage_api", __name__, url_prefix="/api/imessage")

logger = logging.getLogger(__name__)

# Rate limiting: per-number message timestamps
_rate_limit_store: dict[str, list[float]] = defaultdict(list)

RATE_LIMIT_MAX = int(os.getenv("IMESSAGE_RATE_LIMIT_MAX", "10"))
RATE_LIMIT_WINDOW = int(os.getenv("IMESSAGE_RATE_LIMIT_WINDOW", "60"))

MAX_MESSAGE_LENGTH = 2000

SENDBLUE_API_BASE = "https://api.sendblue.co"


def _get_sendblue_headers() -> dict[str, str]:
    return {
        "sb-api-key-id": os.getenv("SENDBLUE_API_KEY", ""),
        "sb-api-secret-key": os.getenv("SENDBLUE_API_SECRET", ""),
        "Content-Type": "application/json",
    }


def _check_rate_limit(phone_number: str) -> bool:
    """Check if a phone number has exceeded the rate limit.

    Returns True if the request is allowed, False if rate-limited.
    """
    now = time.monotonic()
    cutoff = now - RATE_LIMIT_WINDOW

    timestamps = _rate_limit_store[phone_number]
    _rate_limit_store[phone_number] = [t for t in timestamps if t > cutoff]

    if len(_rate_limit_store[phone_number]) >= RATE_LIMIT_MAX:
        return False

    _rate_limit_store[phone_number].append(now)
    return True


async def send_imessage(
    to_number: str,
    content: str,
    reply_to_handle: str | None = None,
) -> dict:
    """Send an iMessage via SendBlue API."""
    from_number = os.getenv("SENDBLUE_FROM_NUMBER", "")

    payload: dict = {
        "number": to_number,
        "from_number": from_number,
        "content": content,
    }
    if reply_to_handle:
        payload["reply_to"] = {"message_handle": reply_to_handle}

    async with httpx.AsyncClient() as client:
        resp = await client.post(
            f"{SENDBLUE_API_BASE}/api/send-message",
            headers=_get_sendblue_headers(),
            json=payload,
            timeout=30,
        )
        resp.raise_for_status()
        return resp.json()


async def send_typing_indicator(
    to_number: str,
    state: str = "start",
    max_duration_ms: int | None = None,
) -> None:
    """Send or stop a typing indicator. Errors are logged and swallowed."""
    from_number = os.getenv("SENDBLUE_FROM_NUMBER", "")
    payload: dict = {
        "number": to_number,
        "from_number": from_number,
        "state": state,
    }
    if max_duration_ms is not None:
        payload["max_duration_ms"] = max_duration_ms

    try:
        async with httpx.AsyncClient() as client:
            await client.post(
                f"{SENDBLUE_API_BASE}/api/send-typing-indicator",
                headers=_get_sendblue_headers(),
                json=payload,
                timeout=10,
            )
    except Exception as e:
        logger.warning(f"Failed to send typing indicator ({state}): {e}")


def validate_sendblue_signature(f):
    """Decorator to validate the SendBlue webhook signing secret."""

    @functools.wraps(f)
    async def decorated_function(*args, **kwargs):
        if os.getenv("SENDBLUE_SIGNATURE_VALIDATION", "true").lower() == "false":
            return await f(*args, **kwargs)

        secret = os.getenv("SENDBLUE_WEBHOOK_SECRET")
        if not secret:
            logger.error("SENDBLUE_WEBHOOK_SECRET not set — rejecting request")
            return jsonify({"error": "Server misconfigured"}), 500

        sig = request.headers.get("sb-signing-secret", "")
        if not hmac.compare_digest(sig, secret):
            logger.warning("Invalid SendBlue signing secret")
            return jsonify({"error": "Unauthorized"}), 403

        return await f(*args, **kwargs)

    return decorated_function


@imessage_blueprint.route("/webhook", methods=["POST"])
@validate_sendblue_signature
async def webhook():
    """Handle incoming iMessages from SendBlue."""
    data = await request.get_json()
    if not data:
        return jsonify({"error": "Invalid payload"}), 400

    from_number = data.get("from_number")
    content = data.get("content") or ""
    media_url = data.get("media_url")
    is_outbound = data.get("is_outbound", False)

    # Threading fields from SendBlue
    inbound_handle = data.get("message_handle")
    reply_to_handle = (data.get("reply_to") or {}).get("message_handle")
    originator_handle = (data.get("thread_originator") or {}).get("message_handle")

    # Ignore outbound messages (our own replies echoed back)
    if is_outbound:
        return jsonify({"status": "ignored"}), 200

    if not from_number:
        return jsonify({"error": "Missing from_number"}), 400

    if not content and not media_url:
        return jsonify({"error": "Missing content and media_url"}), 400

    content = content.strip()
    if not content and not media_url:
        await send_imessage(
            from_number,
            "I received an empty message. Please send some text or an image!",
        )
        return jsonify({"status": "ok"}), 200

    # Rate limiting
    if not _check_rate_limit(from_number):
        logger.warning(f"Rate limit exceeded for {from_number}")
        await send_imessage(
            from_number,
            "You're sending messages too quickly. Please wait a moment and try again.",
        )
        return jsonify({"status": "rate_limited"}), 200

    # Truncate overly long messages
    if len(content) > MAX_MESSAGE_LENGTH:
        content = content[:MAX_MESSAGE_LENGTH]
        logger.info(
            f"Truncated long message from {from_number} to {MAX_MESSAGE_LENGTH} chars"
        )

    logger.info(f"Received iMessage from {from_number}: {content[:100]}")

    # Identify or create user
    user = await User.filter(imessage_number=from_number).first()

    if not user:
        allowed_numbers = os.getenv("ALLOWED_IMESSAGE_NUMBERS", "").split(",")
        if from_number not in allowed_numbers and "*" not in allowed_numbers:
            await send_imessage(
                from_number, "Sorry, you are not authorized to use this service."
            )
            return jsonify({"status": "unauthorized"}), 200

        username = f"im_{from_number.lstrip('+')}"
        try:
            user = await User.create(
                username=username,
                email=f"{username}@imessage.simbarag.local",
                imessage_number=from_number,
                auth_provider="imessage",
            )
            logger.info(f"Created new user for iMessage: {username}")
        except Exception as e:
            logger.error(f"Failed to create user for {from_number}: {e}")
            await send_imessage(
                from_number, "Sorry, something went wrong setting up your account."
            )
            return jsonify({"status": "error"}), 200

    # iMessage is restricted to admins
    if not user.is_admin():
        logger.warning(f"Non-admin user {user.username} attempted iMessage access")
        await send_imessage(from_number, "Sorry, this feature is restricted to admins.")
        return jsonify({"status": "forbidden"}), 200

    # /threads command — list recent iMessage conversations
    if content.strip().lower() == "/threads":
        threads = (
            await Conversation.filter(user=user, channel="imessage")
            .order_by("-updated_at")
            .limit(5)
        )
        if not threads:
            await send_imessage(from_number, "No threads found.")
        else:
            lines = ["Your recent threads (most recent first):"]
            for i, t in enumerate(threads, 1):
                updated = t.updated_at.strftime("%b %-d")
                lines.append(f"{i}. {t.name} ({updated})")
            lines.append("\nReply to any message to continue its thread.")
            await send_imessage(from_number, "\n".join(lines))
        return jsonify({"status": "ok"}), 200

    # Determine conversation based on iMessage reply threading:
    # - Reply to an existing message → continue that conversation
    # - New standalone message → start a fresh conversation
    try:
        if reply_to_handle:
            lookup_handle = originator_handle or reply_to_handle
            parent_msg = await ConversationMessage.get_or_none(
                message_handle=lookup_handle,
                conversation__user=user,
                conversation__channel="imessage",
            )
            if parent_msg:
                conversation = await Conversation.get(id=parent_msg.conversation_id)
                logger.info(
                    f"Continuing conversation {conversation.id} via reply to {lookup_handle!r}"
                )
            else:
                # Handle not found (pre-migration messages) — fall back to most recent
                logger.warning(
                    f"Reply handle {lookup_handle!r} not in DB; falling back to most recent conversation"
                )
                conversation = await get_conversation_for_channel(
                    user=user, channel="imessage"
                )
        else:
            conversation = await Conversation.create(
                name="New iMessage",
                user=user,
                channel="imessage",
            )
            logger.info(
                f"New iMessage conversation {conversation.id} for {user.username}"
            )

    except Exception as e:
        logger.error(f"Failed to get/create conversation for {user.username}: {e}")
        await send_imessage(
            from_number, "Sorry, something went wrong. Please try again later."
        )
        return jsonify({"status": "error"}), 200

    # Save the inbound message now so it's persisted before we ack
    await add_message_to_conversation(
        conversation=conversation,
        message=content,
        speaker="user",
        user=user,
        message_handle=inbound_handle,
    )

    # Fire LLM processing in the background so we can ack SendBlue immediately.
    # SendBlue retries if it doesn't receive a 200 within its timeout (~5s).
    asyncio.ensure_future(
        _process_and_reply(
            from_number=from_number,
            inbound_handle=inbound_handle,
            conversation=conversation,
            user=user,
            content=content,
            media_url=media_url,
        )
    )

    return jsonify({"status": "ok"}), 200


async def _process_and_reply(
    from_number: str,
    inbound_handle: str | None,
    conversation: Conversation,
    user,
    content: str,
    media_url: str | None,
) -> None:
    """Run the LLM agent and send the reply. Called as a background task."""
    # Attach image if present
    image_block = None
    if media_url:
        try:
            from utils.image_process import image_content_block

            async with httpx.AsyncClient() as client:
                resp = await client.get(media_url, timeout=30)
                resp.raise_for_status()
            image_block = image_content_block(resp.content)
        except Exception as e:
            logger.error(f"Failed to load iMessage image: {e}")
            content = f"{content or ''}\n\n[An image was sent but could not be loaded.]".strip()

    await send_typing_indicator(from_number, state="start", max_duration_ms=120_000)

    response_text = "Sorry, I'm having trouble thinking right now."
    try:
        system_prompt = await build_system_prompt_with_memories(str(user.id))
        system_prompt += (
            "\n\nIMPORTANT: You are replying via iMessage. Use plain text only. "
            "Do not use markdown tables, headers, bullet points, bold, italics, "
            "or any other markdown formatting. Write in natural prose."
        )

        if image_block:
            user_content = [
                {"type": "text", "text": content or "The user sent this image."},
                image_block,
            ]
        else:
            user_content = content

        agent_config = {
            "configurable": {
                "user_id": str(user.id),
                "is_admin": user.is_admin(),
                "thread_id": str(conversation.id),
                "system_prompt": system_prompt,
            }
        }
        logger.info(f"Invoking LangChain agent for conversation {conversation.id}")
        response = await main_agent.ainvoke(
            {"messages": [{"role": "user", "content": user_content}]},
            config=agent_config,
        )
        response_text = extract_agent_text(response.get("messages", [])[-1].content)
        await persist_agent_tool_messages(
            conversation, user, response.get("messages", [])
        )

    except Exception as e:
        logger.error(f"Error invoking agent: {e}")
    finally:
        await send_typing_indicator(from_number, state="stop")

    send_resp = await send_imessage(
        from_number,
        strip_markdown(response_text),
        reply_to_handle=inbound_handle,
    )
    outbound_handle = (send_resp or {}).get("message_handle")

    await add_message_to_conversation(
        conversation=conversation,
        message=response_text,
        speaker="simba",
        user=user,
        message_handle=outbound_handle,
    )