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

from __future__ import annotations

from dataclasses import dataclass, field
from enum import StrEnum
from typing import Any


class ContextSource(StrEnum):
    CONTACT = "contact"
    LOCATION = "location"
    REACTION = "reaction"
    POLL = "poll"
    BUTTON = "button"
    FORWARD = "forward"
    MENTION = "mention"
    MEDIA = "media"
    UNKNOWN = "unknown"


class ContextType(StrEnum):
    STRUCTURED_DATA = "structured_data"
    USER_GENERATED = "user_generated"
    SYSTEM = "system"


@dataclass
class StructuredContextEntry:
    label: str
    source: ContextSource = ContextSource.UNKNOWN
    type: ContextType = ContextType.STRUCTURED_DATA
    data: dict[str, Any] = field(default_factory=dict)


@dataclass
class UntrustedStructuredContext:
    entries: list[StructuredContextEntry] = field(default_factory=list)

    @classmethod
    def from_contact(
        cls,
        display_name: str = "",
        vcard: str = "",
        contacts: list[dict[str, str]] | None = None,
    ) -> UntrustedStructuredContext:
        entry = StructuredContextEntry(
            label="contact",
            source=ContextSource.CONTACT,
            data={
                "display_name": display_name,
                "vcard": vcard,
                "contacts": contacts or [],
            },
        )
        return cls(entries=[entry])

    @classmethod
    def from_location(
        cls, latitude: float, longitude: float, name: str = "", address: str = ""
    ) -> UntrustedStructuredContext:
        entry = StructuredContextEntry(
            label="location",
            source=ContextSource.LOCATION,
            data={
                "latitude": latitude,
                "longitude": longitude,
                "name": name,
                "address": address,
            },
        )
        return cls(entries=[entry])

    @classmethod
    def from_reaction(cls, emoji: str, target_msg_id: str, target_jid: str) -> UntrustedStructuredContext:
        entry = StructuredContextEntry(
            label="reaction",
            source=ContextSource.REACTION,
            data={
                "emoji": emoji,
                "target_msg_id": target_msg_id,
                "target_jid": target_jid,
            },
        )
        return cls(entries=[entry])

    @classmethod
    def from_poll(
        cls, poll_name: str, options: list[str], selected_options: list[str] | None = None
    ) -> UntrustedStructuredContext:
        entry = StructuredContextEntry(
            label="poll",
            source=ContextSource.POLL,
            data={
                "name": poll_name,
                "options": options,
                "selected_options": selected_options or [],
            },
        )
        return cls(entries=[entry])

    @classmethod
    def from_button(cls, button_id: str, display_text: str) -> UntrustedStructuredContext:
        entry = StructuredContextEntry(
            label="button",
            source=ContextSource.BUTTON,
            data={
                "button_id": button_id,
                "display_text": display_text,
            },
        )
        return cls(entries=[entry])

    @classmethod
    def from_forward(cls, forwarding_score: int = 0, forwarded_from: str = "") -> UntrustedStructuredContext:
        entry = StructuredContextEntry(
            label="forward",
            source=ContextSource.FORWARD,
            data={
                "is_forwarded": True,
                "forwarding_score": forwarding_score,
                "forwarded_from": forwarded_from,
            },
        )
        return cls(entries=[entry])

    @classmethod
    def from_mentions(cls, mentioned_jids: list[str]) -> UntrustedStructuredContext:
        entry = StructuredContextEntry(
            label="mention",
            source=ContextSource.MENTION,
            data={"mentioned_jids": mentioned_jids},
        )
        return cls(entries=[entry])

    def add_entry(self, entry: StructuredContextEntry) -> None:
        self.entries.append(entry)

    def get_by_label(self, label: str) -> StructuredContextEntry | None:
        for entry in self.entries:
            if entry.label == label:
                return entry
        return None

    def get_by_source(self, source: ContextSource) -> list[StructuredContextEntry]:
        return [e for e in self.entries if e.source == source]

    def to_metadata(self) -> dict[str, Any]:
        result: dict[str, Any] = {}
        for entry in self.entries:
            result[f"structured_{entry.label}"] = {
                "source": entry.source.value,
                "type": entry.type.value,
                "data": entry.data,
            }
        return result

    def is_empty(self) -> bool:
        return len(self.entries) == 0