vk_hackathon/index/cleaning.py

"""Local message cleaning and normalization. No external API calls."""

import json
import re
from typing import Any

_ZERO_WIDTH = re.compile(r"[\u200b\u200c\u200d\ufeff]")
_MULTI_NEWLINE = re.compile(r"\n{3,}")
_MULTI_SPACE = re.compile(r"[ \t]{2,}")


def normalize_unicode(text: str) -> str:
    text = _ZERO_WIDTH.sub("", text)
    text = text.replace("\r\n", "\n").replace("\r", "\n")
    text = _MULTI_SPACE.sub(" ", text)
    text = _MULTI_NEWLINE.sub("\n\n", text)
    return text.strip()


def _safe_normalize(text: str | None) -> str:
    if not text:
        return ""
    return normalize_unicode(str(text))


def parse_file_snippets(raw: str) -> list[dict[str, Any]]:
    if not raw or not raw.strip():
        return []
    try:
        parsed = json.loads(raw)
        if isinstance(parsed, list):
            return parsed
        if isinstance(parsed, dict):
            return [parsed]
        return []
    except (json.JSONDecodeError, ValueError):
        return []


def extract_file_info(snippet: dict[str, Any]) -> dict[str, str]:
    return {
        "name": str(snippet.get("name") or ""),
        "mime": str(snippet.get("mime") or ""),
        "url": str(snippet.get("original_url") or ""),
        "date": str(snippet.get("date_create") or ""),
    }


def normalize_member_event(event: dict[str, Any] | None) -> str:
    if not event:
        return ""
    event_type = str(event.get("type") or "unknown_event")
    members = event.get("members") or []
    if event_type == "addMembers" and members:
        joined = ", ".join(str(m) for m in members)
        return f"[system: {joined} added to chat]"
    payload_str = json.dumps(event, ensure_ascii=False)
    return f"[system: {event_type} {payload_str}]"


def normalize_part(part: dict[str, Any]) -> dict[str, str]:
    """Normalize a single message part by its mediaType."""
    media_type = str(part.get("mediaType") or "text")
    text = _safe_normalize(part.get("text"))

    if media_type == "text":
        return {"type": "text", "text": text}

    if media_type == "quote":
        sender = _safe_normalize(part.get("sn") or part.get("sender_id") or "")
        label = f"[quote from {sender}]" if sender else "[quote]"
        return {"type": "quote", "text": f"{label}: {text}" if text else label}

    if media_type == "forward":
        origin = _safe_normalize(part.get("sn") or "")
        label = f"[forwarded from {origin}]" if origin else "[forwarded]"
        return {"type": "forward", "text": f"{label}: {text}" if text else label}

    label = f"[{media_type}]"
    return {"type": media_type, "text": f"{label}: {text}" if text else label}


class CleanedMessage:
    __slots__ = (
        "id",
        "sender_id",
        "time",
        "thread_sn",
        "text",
        "parts",
        "mentions",
        "member_event_text",
        "file_info",
        "is_system",
        "is_forward",
        "is_quote",
        "is_empty",
    )

    def __init__(
        self,
        id: str,
        sender_id: str,
        time: int,
        thread_sn: str | None,
        text: str,
        parts: list[dict[str, str]],
        mentions: list[str],
        member_event_text: str,
        file_info: list[dict[str, str]],
        is_system: bool,
        is_forward: bool,
        is_quote: bool,
    ):
        self.id = id
        self.sender_id = sender_id
        self.time = time
        self.thread_sn = thread_sn
        self.text = text
        self.parts = parts
        self.mentions = mentions
        self.member_event_text = member_event_text
        self.file_info = file_info
        self.is_system = is_system
        self.is_forward = is_forward
        self.is_quote = is_quote
        self.is_empty = not (text or parts or member_event_text or file_info)


def clean_message(msg: Any) -> CleanedMessage:
    """Extract and normalize all signals from a raw Message object."""
    text = _safe_normalize(msg.text)
    parts = [normalize_part(p) for p in (msg.parts or []) if isinstance(p, dict)]
    parts = [p for p in parts if p.get("text")]
    mentions = [_safe_normalize(m) for m in (msg.mentions or []) if m]
    member_event_text = normalize_member_event(msg.member_event)

    file_info: list[dict[str, str]] = []
    for snippet in parse_file_snippets(msg.file_snippets):
        info = extract_file_info(snippet)
        if any(info.values()):
            file_info.append(info)

    return CleanedMessage(
        id=msg.id,
        sender_id=_safe_normalize(msg.sender_id),
        time=msg.time,
        thread_sn=msg.thread_sn,
        text=text,
        parts=parts,
        mentions=mentions,
        member_event_text=member_event_text,
        file_info=file_info,
        is_system=bool(msg.is_system),
        is_forward=bool(msg.is_forward),
        is_quote=bool(msg.is_quote),
    )