"""
Export an iMessage conversation with one contact to JSON + CSV from your
Mac's local Messages database.

Usage:
    python3 export_conversation.py --contact "+16095551234"
    python3 export_conversation.py --contact 6095551234 --db ~/Desktop/chat_copy.db
    python3 export_conversation.py --contact friend@icloud.com

The contact can be a phone number (any format — last 10 digits are matched)
or an email (matched as a substring against the iMessage handle).

Before running:
  1. Grant Full Disk Access to your terminal (System Settings -> Privacy &
     Security -> Full Disk Access).
  2. Copy your live Messages DB to a safe spot:
     cp ~/Library/Messages/chat.db ~/Desktop/chat_copy.db
"""

import argparse
import csv
import json
import re
import sqlite3
from datetime import datetime, timezone
from pathlib import Path

APPLE_EPOCH = datetime(2001, 1, 1, tzinfo=timezone.utc)


def apple_ts_to_iso(ns):
    if not ns:
        return None
    absolute = APPLE_EPOCH.timestamp() + (ns / 1_000_000_000)
    return datetime.fromtimestamp(absolute, tz=timezone.utc).isoformat()


def decode_attributed_body(blob):
    if not blob:
        return None
    m = re.search(rb"NSString[^+]*\+(.)([^\x86]*)", blob, re.DOTALL)
    if not m:
        return None
    return m.group(2).decode("utf-8", errors="ignore") or None


def normalize_phone_tail(s):
    digits = re.sub(r"\D", "", s)
    return digits[-10:] if len(digits) >= 10 else None


def find_handles(conn, contact):
    tail = normalize_phone_tail(contact)
    if tail:
        return conn.execute(
            """
            SELECT ROWID, id FROM handle
             WHERE REPLACE(REPLACE(REPLACE(REPLACE(id,' ',''),'-',''),'(',''),')','')
                   LIKE ?
            """,
            (f"%{tail}",),
        ).fetchall()
    # Otherwise treat as email substring
    return conn.execute(
        "SELECT ROWID, id FROM handle WHERE id LIKE ?",
        (f"%{contact}%",),
    ).fetchall()


def main():
    ap = argparse.ArgumentParser(description="Export an iMessage conversation.")
    ap.add_argument("--contact", required=True,
                    help="Phone number (any format) or email of the other person.")
    ap.add_argument("--db", default="~/Desktop/chat_copy.db",
                    help="Path to a copy of chat.db (default: ~/Desktop/chat_copy.db)")
    ap.add_argument("--out", default="./out",
                    help="Output directory (default: ./out)")
    args = ap.parse_args()

    db_path = Path(args.db).expanduser()
    out_dir = Path(args.out).expanduser()
    out_dir.mkdir(parents=True, exist_ok=True)

    if not db_path.exists():
        raise SystemExit(
            f"Could not find {db_path}.\n"
            f"Run: cp ~/Library/Messages/chat.db {db_path}"
        )

    conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
    conn.row_factory = sqlite3.Row

    handles = find_handles(conn, args.contact)
    if not handles:
        raise SystemExit(f"No iMessage handle matched '{args.contact}'.")
    print(f"Matched {len(handles)} handle(s): {[h['id'] for h in handles]}")

    handle_ids = [h["ROWID"] for h in handles]
    placeholders = ",".join("?" * len(handle_ids))
    rows = conn.execute(
        f"""
        SELECT
            m.ROWID                AS id,
            m.date                 AS date_ns,
            m.is_from_me           AS is_from_me,
            m.text                 AS text,
            m.attributedBody       AS attributed_body,
            m.service              AS service,
            (SELECT COUNT(*) FROM message_attachment_join maj
              WHERE maj.message_id = m.ROWID) AS attachment_count
          FROM message m
         WHERE m.handle_id IN ({placeholders})
         ORDER BY m.date ASC
        """,
        handle_ids,
    ).fetchall()
    print(f"Pulled {len(rows):,} messages.")

    messages = [{
        "id": r["id"],
        "date": apple_ts_to_iso(r["date_ns"]),
        "from_me": bool(r["is_from_me"]),
        "text": r["text"] or decode_attributed_body(r["attributed_body"]),
        "service": r["service"],
        "attachments": r["attachment_count"],
    } for r in rows]

    json_path = out_dir / "conversation.json"
    json_path.write_text(json.dumps(messages, indent=2, ensure_ascii=False))

    csv_path = out_dir / "conversation.csv"
    with csv_path.open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=["id", "date", "from_me", "text",
                                           "service", "attachments"])
        w.writeheader()
        w.writerows(messages)

    print(f"Wrote {json_path} and {csv_path}")
    if messages:
        print(f"First message: {messages[0]['date']}")
        print(f"Last message:  {messages[-1]['date']}")


if __name__ == "__main__":
    main()