utils.py

"""Shared utilities for the LLM-Imitate pipeline."""

from __future__ import annotations

import json
import re
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import yaml

SYSTEM_MESSAGE_PATTERNS = [
    re.compile(r"^Liked a message$", re.I),
    re.compile(r"^Reacted .+ to your message\s*$", re.I),
    re.compile(r"^(Audio|Video) call ended$", re.I),
    re.compile(r"^You started an (audio call|video chat)$", re.I),
    re.compile(r"^Missed\b", re.I),
    re.compile(r"^Missed this one$", re.I),
    re.compile(r"^You unsent a message$", re.I),
    re.compile(r"^.+ unsent a message$", re.I),
    re.compile(r"^.+ sent an attachment\.?$", re.I),
]

PII_PATTERNS = [
    ("email", re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b")),
    ("phone", re.compile(r"(?<!\d)(?:\+?\d[\d\s().-]{7,}\d)(?!\d)")),
    ("long_number", re.compile(r"\b\d{10,}\b")),
    ("password_mention", re.compile(r"\b(password|passwd|otp|pin)\b", re.I)),
    ("address_hint", re.compile(
        r"\b\d{1,5}\s+\w+\s+(street|st\.|road|rd\.|avenue|ave\.|lane|ln\.|drive|dr\.|block|sector)\b",
        re.I,
    )),
]


def load_config(path: str | Path = "config.yaml") -> dict[str, Any]:
    with open(path, encoding="utf-8") as f:
        return yaml.safe_load(f)


def fix_mojibake(text: str) -> str:
    """Fix Instagram's double-encoded UTF-8 → Latin-1 mojibake."""
    if not text:
        return text
    try:
        return text.encode("latin1").decode("utf8")
    except (UnicodeDecodeError, UnicodeEncodeError):
        return text


def fix_value(value: Any) -> Any:
    if isinstance(value, str):
        return fix_mojibake(value)
    if isinstance(value, list):
        return [fix_value(v) for v in value]
    if isinstance(value, dict):
        return {k: fix_value(v) for k, v in value.items()}
    return value


def ensure_dir(path: str | Path) -> Path:
    p = Path(path)
    p.mkdir(parents=True, exist_ok=True)
    return p


def read_jsonl(path: str | Path) -> list[dict[str, Any]]:
    records: list[dict[str, Any]] = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                records.append(json.loads(line))
    return records


def write_jsonl(path: str | Path, records: list[dict[str, Any]]) -> None:
    with open(path, "w", encoding="utf-8") as f:
        for record in records:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")


def is_system_message(content: str) -> bool:
    text = content.strip()
    if not text:
        return False
    return any(p.search(text) for p in SYSTEM_MESSAGE_PATTERNS)


def has_media_only(msg: dict[str, Any]) -> bool:
    content = (msg.get("content") or "").strip()
    if content:
        return False
    return bool(msg.get("photos") or msg.get("videos") or msg.get("audio_files"))


def extract_message_text(msg: dict[str, Any], include_shares: bool) -> str | None:
    content = (msg.get("content") or "").strip()
    share = msg.get("share") or {}
    share_text = (share.get("share_text") or share.get("link") or "").strip()

    if content and share_text and content == share_text:
        text = content
    elif content:
        text = content
    elif include_shares and share_text:
        text = share_text
    else:
        return None

    text = text.strip()
    if not text or is_system_message(text):
        return None
    return text


def group_into_turns(
    messages: list[dict[str, Any]],
    gap_seconds: int = 60,
) -> list[dict[str, Any]]:
    """Merge consecutive messages from the same sender within gap_seconds."""
    if not messages:
        return []

    turns: list[dict[str, Any]] = []
    current: dict[str, Any] | None = None

    for msg in messages:
        sender = msg["sender"]
        text = msg["text"]
        ts = msg["timestamp"]

        if current is None:
            current = {
                "sender": sender,
                "text": text,
                "timestamp": ts,
                "timestamp_end": ts,
                "message_count": 1,
            }
            continue

        same_sender = current["sender"] == sender
        gap_ms = ts - current["timestamp_end"]
        within_gap = 0 <= gap_ms <= gap_seconds * 1000

        if same_sender and within_gap:
            current["text"] = f"{current['text']}\n{text}"
            current["timestamp_end"] = ts
            current["message_count"] += 1
        else:
            turns.append(current)
            current = {
                "sender": sender,
                "text": text,
                "timestamp": ts,
                "timestamp_end": ts,
                "message_count": 1,
            }

    if current is not None:
        turns.append(current)

    return turns


def role_for_sender(sender: str, my_sender: str, target_sender: str) -> str:
    if sender == target_sender:
        return "them"
    if sender == my_sender:
        return "me"
    return "other"


def chat_role(sender: str, my_sender: str, target_sender: str) -> str:
    """Map to HF chat roles: me=user, them=assistant."""
    return "assistant" if sender == target_sender else "user"


def timestamp_to_period(timestamp_ms: int) -> str:
    dt = datetime.fromtimestamp(timestamp_ms / 1000, tz=timezone.utc)
    return dt.strftime("%Y-%m")


def detect_pii(text: str) -> list[str]:
    flags: list[str] = []
    for label, pattern in PII_PATTERNS:
        if pattern.search(text):
            flags.append(label)
    return flags


def stratified_split(
    records: list[dict[str, Any]],
    train_ratio: float,
    period_key: str = "period",
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    """Split each time bucket independently to preserve era representation."""
    import random

    buckets: dict[str, list[dict[str, Any]]] = {}
    for record in records:
        buckets.setdefault(record[period_key], []).append(record)

    train: list[dict[str, Any]] = []
    val: list[dict[str, Any]] = []

    for bucket_records in buckets.values():
        shuffled = bucket_records.copy()
        random.seed(42)
        random.shuffle(shuffled)
        split_idx = max(1, int(len(shuffled) * train_ratio)) if len(shuffled) > 1 else len(shuffled)
        if len(shuffled) == 1:
            train.extend(shuffled)
        else:
            train.extend(shuffled[:split_idx])
            val.extend(shuffled[split_idx:])

    return train, val


def summarize_persona_for_prompt(profile: dict[str, Any], max_chars: int = 1200) -> str:
    """Condense persona_profile.json for system prompt injection."""
    lines: list[str] = []
    msg = profile.get("message_stats", {})
    lines.append(
        f"Avg length: {msg.get('avg_chars', '?')} chars, "
        f"{msg.get('avg_words', '?')} words. "
        f"Often sends {msg.get('avg_messages_per_turn', '?')} msgs per thought."
    )

    punct = profile.get("punctuation_habits", {})
    habit_labels = {
        "lowercase_only_pct": "mostly lowercase",
        "no_terminal_punct_pct": "skips ending punctuation",
        "ellipsis_heavy_pct": "uses ...",
        "exclamation_heavy_pct": "uses !!",
        "question_heavy_pct": "uses ??",
        "all_caps_words_pct": "ALL CAPS words",
    }
    quirks = [
        f"{habit_labels[k]} ({v}%)"
        for k, v in punct.items()
        if k in habit_labels and isinstance(v, (int, float)) and v >= 30
    ]
    if quirks:
        lines.append("Punctuation: " + ", ".join(quirks[:6]) + ".")

    top_emojis = profile.get("top_emojis", [])[:8]
    if top_emojis:
        em = ", ".join(f"{e['emoji']}({e['count']})" for e in top_emojis)
        lines.append(f"Top emojis: {em}.")

    top_words = profile.get("top_words", [])[:12]
    if top_words:
        words = ", ".join(w["word"] for w in top_words)
        lines.append(f"Common words: {words}.")

    slang = profile.get("slang_abbreviations", [])[:10]
    if slang:
        lines.append("Slang: " + ", ".join(s["phrase"] for s in slang) + ".")

    openers = profile.get("common_openers", [])[:5]
    if openers:
        lines.append("Openers: " + ", ".join(f"\"{o['phrase']}\"" for o in openers) + ".")

    text = "\n".join(lines)
    return text[:max_chars]