parse_export.py

#!/usr/bin/env python3
"""STEP 1 - Parse the Instagram export into conversation.jsonl.

Keeps real text messages and turns media/shares into short placeholders ([reel], [photo]...).
Drops likes, reactions, attachment notices, calls, unsends and other system messages.
"""

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

from utils import ensure_dir, fix_value, load_config, write_jsonl

SYSTEM_PATTERNS = [
    ("like", re.compile(r"^Liked (a|an|your) .+$", re.I | re.S)),
    ("reaction", re.compile(r"^Reacted .+ to (your|a) message\s*$", re.I)),
    ("unsent", re.compile(r"^.* unsent a message\.?$", re.I)),
    ("call", re.compile(r"^(.* )?(audio|video) (call|chat)\b.*$|^Missed\b.*$|^.* started an? .*(call|video chat).*$", re.I)),
    ("nickname", re.compile(r"^.* (set|changed) the nickname .*$", re.I)),
    ("collection", re.compile(r"^.* added to a collection.*$", re.I | re.S)),
    ("group_event", re.compile(r"^.* (named the group|changed the group|joined the group|left the group|added .+ to the group|changed the theme|set the theme|turned (on|off) vanish mode).*$", re.I)),
    ("music_invite", re.compile(r"^(Let.s listen and choose the music together|https://open\.spotify\.com/socialsession/).*$", re.I | re.S)),
]
ATTACHMENT_RE = re.compile(r"^.*sent an attachment\.?$", re.I | re.S)
EDITED_RE = re.compile(r"\s*\(edited\)\s*$")  # export artifact appended to edited messages


def system_kind(content: str) -> str | None:
    for name, pattern in SYSTEM_PATTERNS:
        if pattern.match(content):
            return name
    return None


def share_placeholder(share: dict[str, Any]) -> str:
    link = share.get("link") or ""
    if share.get("profile_share_username"):
        return "[profile]"
    if "/reel" in link:
        return "[reel]"
    if "/stories/" in link:
        return "[story]"
    if "/p/" in link or "/tv/" in link:
        return "[post]"
    return "[link]"


def media_placeholder(raw: dict[str, Any]) -> str | None:
    if raw.get("share"):
        return share_placeholder(raw["share"])
    if raw.get("photos"):
        return "[photo]"
    if raw.get("videos"):
        return "[video]"
    if raw.get("audio_files"):
        return "[voice note]"
    if raw.get("files"):
        return "[file]"
    return None


def classify(raw: dict[str, Any]) -> tuple[str, str | None]:
    """-> (category, text). category is 'text', 'media' or a drop reason."""
    content = (raw.get("content") or "").strip()
    if content:
        sys_kind = system_kind(content)
        if sys_kind:
            return sys_kind, None
        if not ATTACHMENT_RE.match(content):
            return "text", EDITED_RE.sub("", content)
    placeholder = media_placeholder(raw)
    if placeholder:
        return "media", placeholder
    return ("attachment_notice" if content else "empty"), None


def discover_message_files(root: Path) -> list[Path]:
    files = list(root.rglob("message_*.json"))
    if not files:
        raise FileNotFoundError(f"No message_*.json files found under {root}")
    return sorted(files, key=lambda p: (str(p.parent), int(re.search(r"(\d+)", p.stem).group(1))))


def load_messages(root: Path) -> list[dict[str, Any]]:
    out: list[dict[str, Any]] = []
    for path in discover_message_files(root):
        with open(path, encoding="utf-8") as f:
            data = fix_value(json.load(f))
        out.extend(data.get("messages", []))
    out.sort(key=lambda m: m.get("timestamp_ms", 0))
    return out


def parse_messages(
    raw_messages: list[dict[str, Any]], my_sender: str, target_sender: str
) -> tuple[list[dict[str, Any]], Counter]:
    kept: list[dict[str, Any]] = []
    stats: Counter = Counter()
    for raw in raw_messages:
        sender = (raw.get("sender_name") or "").strip()
        role = "them" if sender == target_sender else "me" if sender == my_sender else None
        if role is None:
            stats["drop:other_sender"] += 1
            continue
        category, text = classify(raw)
        if text is None:
            stats[f"drop:{category}"] += 1
            continue
        stats[f"keep:{category}"] += 1
        kept.append({"ts": raw["timestamp_ms"], "role": role, "text": text, "media": category == "media"})
    return kept, stats


def main() -> None:
    cfg = load_config()
    output_dir = ensure_dir(cfg["output_dir"])
    raw = load_messages(Path(cfg["thread_path"]))
    kept, stats = parse_messages(raw, cfg["my_sender"], cfg["target_sender"])
    write_jsonl(output_dir / "conversation.jsonl", kept)

    fmt = lambda ts: datetime.fromtimestamp(ts / 1000, tz=timezone.utc).strftime("%Y-%m-%d")
    print(f"Raw messages: {len(raw):,}  ->  kept {len(kept):,}  ({fmt(kept[0]['ts'])} to {fmt(kept[-1]['ts'])})")
    for key, n in sorted(stats.items()):
        print(f"  {key:26s}{n:>8,}")
    by_role = Counter(m["role"] for m in kept if not m["media"])
    print(f"Text messages: them={by_role['them']:,} me={by_role['me']:,}")
    if not by_role["them"]:
        print(f"WARNING: no text messages from target_sender '{cfg['target_sender']}'")


if __name__ == "__main__":
    main()