parse_export.py

#!/usr/bin/env python3
"""
STEP 1 — Parse Instagram DM export → conversation.jsonl
"""

from __future__ import annotations

import json
import re
from pathlib import Path

from tqdm import tqdm

from utils import (
    ensure_dir,
    extract_message_text,
    fix_value,
    has_media_only,
    load_config,
    write_jsonl,
)


def discover_message_files(thread_path: Path) -> list[Path]:
    files = list(thread_path.glob("message_*.json"))
    if not files:
        raise FileNotFoundError(f"No message_*.json files found in {thread_path}")

    def sort_key(path: Path) -> int:
        match = re.search(r"message_(\d+)\.json$", path.name)
        return int(match.group(1)) if match else 0

    return sorted(files, key=sort_key)


def load_thread_messages(thread_path: Path) -> list[dict]:
    all_messages: list[dict] = []

    for path in discover_message_files(thread_path):
        with open(path, encoding="utf-8") as f:
            data = fix_value(json.load(f))

        participants = [p.get("name", "") for p in data.get("participants", [])]
        for raw in data.get("messages", []):
            raw["_source_file"] = path.name
            raw["_participants"] = participants
            all_messages.append(raw)

    all_messages.sort(key=lambda m: m.get("timestamp_ms", 0))
    return all_messages


def parse_messages(raw_messages: list[dict], include_shares: bool) -> tuple[list[dict], dict]:
    cleaned: list[dict] = []
    stats = {
        "total_raw": len(raw_messages),
        "skipped_system": 0,
        "skipped_media_only": 0,
        "skipped_empty": 0,
        "kept": 0,
    }

    for raw in raw_messages:
        if has_media_only(raw):
            stats["skipped_media_only"] += 1
            continue

        text = extract_message_text(raw, include_shares=include_shares)
        if text is None:
            if (raw.get("content") or "").strip():
                stats["skipped_system"] += 1
            else:
                stats["skipped_empty"] += 1
            continue

        sender = fix_value(raw.get("sender_name", "")).strip()
        cleaned.append(
            {
                "sender": sender,
                "text": text,
                "timestamp": raw.get("timestamp_ms", 0),
            }
        )
        stats["kept"] += 1

    return cleaned, stats


def main() -> None:
    cfg = load_config()
    thread_path = Path(cfg["thread_path"])
    output_dir = ensure_dir(cfg["output_dir"])
    output_path = output_dir / "conversation.jsonl"

    print(f"Loading Instagram export from {thread_path} ...")
    raw_messages = load_thread_messages(thread_path)
    cleaned, stats = parse_messages(raw_messages, include_shares=cfg.get("include_shares", False))

    write_jsonl(output_path, cleaned)

    senders: dict[str, int] = {}
    for msg in cleaned:
        senders[msg["sender"]] = senders.get(msg["sender"], 0) + 1

    if cleaned:
        ts_start = min(m["timestamp"] for m in cleaned)
        ts_end = max(m["timestamp"] for m in cleaned)
        from datetime import datetime, timezone

        d0 = datetime.fromtimestamp(ts_start / 1000, tz=timezone.utc).strftime("%Y-%m-%d")
        d1 = datetime.fromtimestamp(ts_end / 1000, tz=timezone.utc).strftime("%Y-%m-%d")
        date_range = f"{d0} to {d1}"
    else:
        date_range = "n/a"

    print(f"\nParsed {stats['kept']:,} messages -> {output_path}")
    print(f"  Raw messages:      {stats['total_raw']:,}")
    print(f"  Skipped (system):  {stats['skipped_system']:,}")
    print(f"  Skipped (media):   {stats['skipped_media_only']:,}")
    print(f"  Skipped (empty):   {stats['skipped_empty']:,}")
    print(f"  Date range:        {date_range}")
    print(f"  Senders:")
    for name, count in sorted(senders.items(), key=lambda x: -x[1]):
        print(f"    {name}: {count:,}")

    target = cfg.get("target_sender", "")
    mine = cfg.get("my_sender", "")
    if target and target not in senders:
        print(f"\nWARNING: target_sender '{target}' not found in parsed messages.")
    if mine and mine not in senders:
        print(f"WARNING: my_sender '{mine}' not found in parsed messages.")


if __name__ == "__main__":
    main()