ForcePilot/backend/package/yuxi/channels/adapters/whatsapp/format.py

from __future__ import annotations

import time
from datetime import datetime
from typing import Any

from yuxi.channels.models import (
    Attachment,
    ChannelIdentity,
    ChannelMessage,
    ChannelType,
    ChatType,
    EventType,
    MessageType,
)

from .structured_context import ContextSource, StructuredContextEntry, UntrustedStructuredContext

_PROTOCOL_REVOKE = 0


def normalize_inbound(raw_payload: dict[str, Any], channel_id: str) -> ChannelMessage:
    key = raw_payload.get("key", {})
    msg = raw_payload.get("message", {})

    remote_jid = key.get("remoteJid", "unknown")
    msg_id = key.get("id")
    from_me = key.get("fromMe", False)

    if "protocolMessage" in msg:
        return _normalize_protocol(msg["protocolMessage"], remote_jid, channel_id)

    if "pollUpdateMessage" in msg:
        return _normalize_poll_vote(msg["pollUpdateMessage"], raw_payload, channel_id)

    if from_me:
        content_type, content, attachments = _extract_content(msg, msg_id)
        msg_type = _map_message_type(content_type)
        chat_type = _jid_to_chat_type(remote_jid)
        sender_id = _extract_sender_number(remote_jid)

        return ChannelMessage(
            identity=ChannelIdentity(
                channel_id=channel_id,
                channel_type=ChannelType.WHATSAPP,
                channel_user_id=sender_id,
                channel_chat_id=remote_jid,
                channel_message_id=msg_id,
            ),
            event_type=EventType.MESSAGE_RECEIVED,
            message_type=msg_type,
            chat_type=chat_type,
            content="",
            attachments=[],
            metadata={
                "from_me": True,
                "push_name": raw_payload.get("pushName", ""),
                "raw_type": content_type,
                "remote_jid": remote_jid,
            },
        )

    content_type, content, attachments = _extract_content(msg, msg_id)

    msg_type = _map_message_type(content_type)
    chat_type = _jid_to_chat_type(remote_jid)
    sender_id = _extract_sender_number(remote_jid)

    metadata = {
        "push_name": raw_payload.get("pushName", ""),
        "broadcast": raw_payload.get("broadcast", False),
        "raw_type": content_type,
        "remote_jid": remote_jid,
        **(_extract_forward_info(msg) or {}),
    }

    structured = _build_structured_context(msg, content_type, msg_id)
    if not structured.is_empty():
        metadata.update(structured.to_metadata())

    if content_type == "reaction":
        reaction = msg.get("reactionMessage", {})
        target_key = reaction.get("key", {})
        if target_key:
            metadata["reaction_target_msg_id"] = target_key.get("id")
            metadata["reaction_target_jid"] = target_key.get("remoteJid")

    return ChannelMessage(
        identity=ChannelIdentity(
            channel_id=channel_id,
            channel_type=ChannelType.WHATSAPP,
            channel_user_id=sender_id,
            channel_chat_id=remote_jid,
            channel_message_id=msg_id,
        ),
        event_type=EventType.MESSAGE_RECEIVED,
        message_type=msg_type,
        chat_type=chat_type,
        content=content,
        attachments=attachments,
        timestamp=datetime.fromtimestamp(raw_payload.get("messageTimestamp", time.time())),
        metadata=metadata,
    )


def _normalize_protocol(proto: dict, remote_jid: str, channel_id: str) -> ChannelMessage:
    proto_type = proto.get("type", -1)
    revoked_key = proto.get("key", {})
    revoked_msg_id = revoked_key.get("id", "")

    if proto_type == _PROTOCOL_REVOKE:
        return ChannelMessage(
            identity=ChannelIdentity(
                channel_id=channel_id,
                channel_type=ChannelType.WHATSAPP,
                channel_user_id=_extract_sender_number(remote_jid),
                channel_chat_id=remote_jid,
                channel_message_id=revoked_msg_id,
            ),
            event_type=EventType.MESSAGE_DELETED,
            message_type=MessageType.TEXT,
            chat_type=_jid_to_chat_type(remote_jid),
            content="",
            metadata={"protocol_type": "revoke", "remote_jid": remote_jid},
        )

    return ChannelMessage(
        identity=ChannelIdentity(
            channel_id=channel_id,
            channel_type=ChannelType.WHATSAPP,
            channel_user_id=_extract_sender_number(remote_jid),
            channel_chat_id=remote_jid,
            channel_message_id=revoked_msg_id,
        ),
        event_type=EventType.SYSTEM_EVENT,
        message_type=MessageType.TEXT,
        chat_type=_jid_to_chat_type(remote_jid),
        content="",
        metadata={"protocol_type": proto_type, "remote_jid": remote_jid},
    )


def _normalize_poll_vote(poll_update: dict, raw_payload: dict, channel_id: str) -> ChannelMessage:
    key = raw_payload.get("key", {})
    remote_jid = key.get("remoteJid", "unknown")
    msg_id = key.get("id")

    poll_creation_key = poll_update.get("pollCreationMessageKey", {})
    vote_info = poll_update.get("vote", {})

    selected_options = []
    if isinstance(vote_info, dict):
        selected = vote_info.get("selectedOptions", [])
        if isinstance(selected, list):
            selected_options = selected
        elif isinstance(vote_info, list):
            selected_options = vote_info

    option_names = [o.get("name", str(o)) if isinstance(o, dict) else str(o) for o in selected_options]
    vote_text = f"[Poll Vote] {' | '.join(option_names)}" if option_names else "[Poll Vote]"

    return ChannelMessage(
        identity=ChannelIdentity(
            channel_id=channel_id,
            channel_type=ChannelType.WHATSAPP,
            channel_user_id=_extract_sender_number(remote_jid),
            channel_chat_id=remote_jid,
            channel_message_id=msg_id,
        ),
        event_type=EventType.MESSAGE_RECEIVED,
        message_type=MessageType.TEXT,
        chat_type=_jid_to_chat_type(remote_jid),
        content=vote_text,
        metadata={
            "raw_type": "poll_vote",
            "remote_jid": remote_jid,
            "poll_msg_id": poll_creation_key.get("id"),
            "poll_remote_jid": poll_creation_key.get("remoteJid"),
            "selected_options": selected_options,
        },
    )


def _extract_content(msg: dict, msg_id: str | None) -> tuple[str, str, list[Attachment]]:
    if "conversation" in msg:
        return "text", msg["conversation"], []

    if "extendedTextMessage" in msg:
        ext = msg["extendedTextMessage"]
        text = ext.get("text", "")
        context_info = ext.get("contextInfo", {})

        if "mentionedJid" in context_info:
            mentions = context_info.get("mentionedJid", []) or []
            return "text", text, _build_mention_attachments(text, mentions, msg_id)

        return "text", text, []

    if "imageMessage" in msg:
        img = msg["imageMessage"]
        return (
            "image",
            img.get("caption", ""),
            [
                Attachment(
                    type="image",
                    url=img.get("url"),
                    mime_type=img.get("mimetype"),
                    file_id=msg_id,
                )
            ],
        )

    if "videoMessage" in msg:
        vid = msg["videoMessage"]
        return (
            "video",
            vid.get("caption", ""),
            [
                Attachment(
                    type="video",
                    url=vid.get("url"),
                    mime_type=vid.get("mimetype"),
                    file_id=msg_id,
                )
            ],
        )

    if "audioMessage" in msg:
        aud = msg["audioMessage"]
        return (
            "audio",
            "",
            [
                Attachment(
                    type="audio",
                    url=aud.get("url"),
                    mime_type=aud.get("mimetype"),
                    file_id=msg_id,
                )
            ],
        )

    if "documentMessage" in msg:
        doc = msg["documentMessage"]
        return (
            "file",
            doc.get("caption", ""),
            [
                Attachment(
                    type="file",
                    url=doc.get("url"),
                    filename=doc.get("filename"),
                    mime_type=doc.get("mimetype"),
                    file_id=msg_id,
                )
            ],
        )

    if "reactionMessage" in msg:
        reaction = msg["reactionMessage"]
        return (
            "reaction",
            reaction.get("text", ""),
            [],
        )

    if "buttonsResponseMessage" in msg:
        br = msg["buttonsResponseMessage"]
        return "text", br.get("selectedDisplayText", br.get("selectedButtonId", "")), []

    if "listResponseMessage" in msg:
        lr = msg["listResponseMessage"]
        return "text", lr.get("title", lr.get("description", "")), []

    if "templateButtonReplyMessage" in msg:
        tr = msg["templateButtonReplyMessage"]
        return "text", tr.get("selectedDisplayText", tr.get("selectedId", "")), []

    if "pollCreationMessage" in msg:
        poll = msg["pollCreationMessage"]
        name = poll.get("name", "")
        options = [o.get("optionName", "") for o in poll.get("options", [])]
        return "text", f"[Poll] {name}: {' | '.join(options)}", []

    if "locationMessage" in msg:
        loc = msg["locationMessage"]
        return "location", f"Location: {loc.get('degreesLatitude')},{loc.get('degreesLongitude')}", []

    if "stickerMessage" in msg:
        stk = msg["stickerMessage"]
        return (
            "sticker",
            "",
            [
                Attachment(
                    type="sticker",
                    url=stk.get("url"),
                    file_id=msg_id,
                )
            ],
        )

    if "contactMessage" in msg:
        ct = msg["contactMessage"]
        display = ct.get("displayName", "")
        vcard = ct.get("vcard", "")
        content_parts = [display] if display else []
        if vcard:
            content_parts.append(f"[vCard: {len(vcard)} bytes]")
        return "card", " | ".join(content_parts) if content_parts else "[Contact]", []

    if "contactsArrayMessage" in msg:
        cam = msg["contactsArrayMessage"]
        names = [c.get("displayName", "") for c in cam.get("contacts", [])]
        return "card", ", ".join(filter(None, names)) or f"[{len(cam.get('contacts', []))} contacts]", []

    return "text", str(msg)[:500], []


def _map_message_type(content_type: str) -> MessageType:
    mapping: dict[str, MessageType] = {
        "text": MessageType.TEXT,
        "image": MessageType.IMAGE,
        "video": MessageType.VIDEO,
        "audio": MessageType.AUDIO,
        "file": MessageType.FILE,
        "location": MessageType.LOCATION,
        "sticker": MessageType.STICKER,
        "card": MessageType.CARD,
    }
    return mapping.get(content_type, MessageType.TEXT)


def _build_mention_attachments(text: str, mentioned_jids: list[str], msg_id: str | None) -> list[Attachment]:
    attachments: list[Attachment] = []
    for mj in mentioned_jids:
        phone = mj.split("@")[0]
        attachments.append(
            Attachment(
                type="mention",
                file_id=msg_id,
                filename=phone,
                metadata={"jid": mj, "phone": phone},
            )
        )
    return attachments


def _extract_forward_info(msg: dict) -> dict[str, Any] | None:
    for key in ("extendedTextMessage", "imageMessage", "videoMessage", "audioMessage", "documentMessage"):
        content = msg.get(key, {})
        if not content:
            continue
        context_info = content.get("contextInfo", {})
        if context_info.get("isForwarded") or context_info.get("forwardedNewsletterMessageInfo"):
            return {
                "is_forwarded": True,
                "forwarding_score": context_info.get("forwardingScore", 0),
                "forwarded_from": context_info.get("forwardedNewsletterMessageInfo", {}).get("newsletterJid", ""),
            }
    return None


def _jid_to_chat_type(jid: str) -> ChatType:
    if "@g.us" in jid:
        return ChatType.GROUP
    if "@broadcast" in jid:
        return ChatType.GUILD_CHANNEL
    return ChatType.DIRECT


def _extract_sender_number(jid: str) -> str:
    return jid.split("@")[0]


def _build_structured_context(msg: dict, content_type: str, msg_id: str | None) -> UntrustedStructuredContext:
    context = UntrustedStructuredContext()

    if content_type == "location" and "locationMessage" in msg:
        loc = msg["locationMessage"]
        context.add_entry(
            StructuredContextEntry(
                label="location",
                source=ContextSource.LOCATION,
                data={
                    "latitude": loc.get("degreesLatitude"),
                    "longitude": loc.get("degreesLongitude"),
                    "name": loc.get("name", ""),
                    "address": loc.get("address", ""),
                },
            )
        )

    if content_type == "reaction" and "reactionMessage" in msg:
        reaction = msg["reactionMessage"]
        target_key = reaction.get("key", {})
        context.add_entry(
            StructuredContextEntry(
                label="reaction",
                source=ContextSource.REACTION,
                data={
                    "emoji": reaction.get("text", ""),
                    "target_msg_id": target_key.get("id"),
                    "target_jid": target_key.get("remoteJid"),
                },
            )
        )

    if content_type == "card" and "contactMessage" in msg:
        ct = msg["contactMessage"]
        vcard = ct.get("vcard", "")
        display = ct.get("displayName", "")
        context.add_entry(
            StructuredContextEntry(
                label="contact",
                source=ContextSource.CONTACT,
                data={
                    "display_name": display,
                    "vcard": vcard,
                    "contacts": [],
                },
            )
        )

    if content_type == "card" and "contactsArrayMessage" in msg:
        cam = msg["contactsArrayMessage"]
        contacts_list = []
        for c in cam.get("contacts", []):
            contacts_list.append({"displayName": c.get("displayName", ""), "vcard": c.get("vcard", "")})
        context.add_entry(
            StructuredContextEntry(
                label="contacts_array",
                source=ContextSource.CONTACT,
                data={"contacts": contacts_list},
            )
        )

    if content_type == "text" and "pollCreationMessage" in msg:
        poll = msg["pollCreationMessage"]
        options = [o.get("optionName", "") for o in poll.get("options", [])]
        context.add_entry(
            StructuredContextEntry(
                label="poll",
                source=ContextSource.POLL,
                data={
                    "name": poll.get("name", ""),
                    "options": options,
                },
            )
        )

    if content_type == "text" and (
        "buttonsResponseMessage" in msg or "templateButtonReplyMessage" in msg or "listResponseMessage" in msg
    ):
        btn = msg.get("buttonsResponseMessage") or msg.get("templateButtonReplyMessage") or {}
        lst = msg.get("listResponseMessage") or {}
        button_id = btn.get("selectedButtonId", btn.get("selectedId", ""))
        display = btn.get("selectedDisplayText", lst.get("title", lst.get("description", "")))
        if button_id or display:
            context.add_entry(
                StructuredContextEntry(
                    label="button",
                    source=ContextSource.BUTTON,
                    data={
                        "button_id": button_id,
                        "display_text": display,
                    },
                )
            )

    forward_info = _extract_forward_info(msg)
    if forward_info and forward_info.get("is_forwarded"):
        context.add_entry(
            StructuredContextEntry(
                label="forward",
                source=ContextSource.FORWARD,
                data={
                    "is_forwarded": True,
                    "forwarding_score": forward_info.get("forwarding_score", 0),
                    "forwarded_from": forward_info.get("forwarded_from", ""),
                },
            )
        )

    for msg_key in ("extendedTextMessage",):
        ext = msg.get(msg_key, {})
        context_info = ext.get("contextInfo", {})
        mentioned = context_info.get("mentionedJid", []) or []
        if mentioned:
            context.add_entry(
                StructuredContextEntry(
                    label="mention",
                    source=ContextSource.MENTION,
                    data={"mentioned_jids": mentioned},
                )
            )
            break

    return context


def format_outbound(response) -> dict[str, Any]:
    result: dict[str, Any] = {
        "jid": response.identity.channel_chat_id,
        "content": response.content,
    }
    if response.reply_to_message_id:
        result["reply_to"] = response.reply_to_message_id
    if response.message_type == MessageType.IMAGE and response.attachments:
        result["media_type"] = "image"
        result["media_path"] = response.attachments[0].url or response.attachments[0].file_id
    return result