#!/usr/bin/python3
# citations.py: link each sentence of a conversation's summary back to where it was said in the
# transcript, the macOS equivalent of the iOS "tap a summary line -> jump to the transcript" feature.
# It uses the SAME deterministic matcher the iOS app uses: content-word overlap between the summary
# sentence and each transcript block, accepted at >= 0.30 here (the fraction of the summary sentence's
# content words that appear in the block; see best_block for why the Mac uses 0.30, not iOS's 0.45).
# Pure stdlib, no model, no network.
#
#   citations.py --rec-dir DIR --conv-key FIRST_CLIP_FILE
#
# Output JSON: {ok, title, who_with, time_range,
#               summary_sentences:[{text, block: <int|null>}],   # block = index into blocks[], or null
#               blocks:[ "transcript sentence", ... ],
#               block_meta:[{time: "h:mmam"|"", gap: "N min quiet"|null}, ...]}  # index-aligned to blocks
# Pure stdlib.
import sys, os, json, re, datetime

HOME = os.path.expanduser("~")
DEFAULT_REC = os.path.join(HOME, "Recordings/Epiphani")

# The iOS ConversationSummarizer stopword set (function words, fillers, contractions without
# apostrophes). Kept verbatim so the Mac matcher accepts/rejects exactly what iOS does.
STOPWORDS = {
    "the", "and", "for", "are", "but", "not", "you", "all", "any", "can", "had", "her", "was",
    "one", "our", "out", "has", "him", "his", "how", "man", "new", "now", "old", "see", "two",
    "way", "who", "boy", "did", "its", "let", "put", "say", "she", "too", "use", "that", "with",
    "have", "this", "will", "your", "from", "they", "know", "want", "been", "good", "much", "some",
    "time", "very", "when", "come", "here", "just", "like", "long", "make", "many", "more", "only",
    "over", "such", "take", "than", "them", "then", "well", "were", "what", "would", "there",
    "their", "about", "could", "other", "these", "thing", "think", "going", "really", "yeah",
    "okay", "gonna", "kind", "sort", "stuff", "maybe", "right", "actually", "basically", "literally",
    "mean", "guess", "something", "anything", "everything", "someone", "anyone", "everyone",
    "because", "where", "which", "while", "should", "being", "doing", "into", "onto", "yours",
    "mine", "ours", "theirs", "uh", "um", "hmm", "huh", "gotta", "wanna", "dont", "didnt", "thats",
    "youre", "weve", "ive", "youve", "theyre", "lets", "cant", "wont", "shes", "hes", "before",
    "after", "still", "also", "even", "around", "though", "another", "until",
}


def content_tokens(s):
    """iOS contentTokens: lowercase, split on non-alphanumeric, keep tokens > 2 chars that are not
    stopwords, as a set. \\w keeps unicode letters/digits (and '_', which our >2 + stopword filter
    leaves harmless), so non-ASCII content words survive."""
    toks = re.split(r"[^\w]+", s.lower(), flags=re.UNICODE)
    return {t for t in toks if len(t) > 2 and t not in STOPWORDS}


def best_block(sentence, block_tokens, threshold=0.30, min_shared=2):
    """Index of the block with the highest overlap, or None. overlap = |sentence ∩ block| / |sentence|
    (directional, NOT Jaccard). Guard: a sentence with < 2 content words never matches (the iOS guard).

    iOS uses 0.45 because its summaries are EXTRACTIVE (they reuse the transcript's own words, so a
    real citation overlaps heavily). The Mac's summaries are ABSTRACTIVE (GPT paraphrases), so the
    same true citation overlaps less: verified on real data, a correct match ("downstairs at the
    Sutton") lands at 0.33, and the abstractive whole-conversation overview line correctly stays
    unmatched at 0.15. So the threshold is 0.30 here, with a precision floor of >= 2 shared content
    words so a short sentence cannot false-match on a single coincidental word."""
    line = content_tokens(sentence)
    if len(line) < 2:
        return None
    best_i, best, best_shared = None, -1.0, 0
    for i, bt in enumerate(block_tokens):
        if not bt:
            continue
        shared = len(line & bt)
        ov = shared / len(line)
        if ov > best:
            best, best_i, best_shared = ov, i, shared
    return best_i if (best_i is not None and best >= threshold and best_shared >= min_shared) else None


def split_sentences(text):
    """Split prose into sentences on ., !, ? boundaries. Robust to ASR transcripts: if there is no
    terminal punctuation the whole string is one block."""
    text = (text or "").strip()
    if not text:
        return []
    parts = re.split(r"(?<=[.!?])\s+", text)
    return [p.strip() for p in parts if p.strip()]


def load_diary(rec_dir):
    try:
        with open(os.path.join(rec_dir, "diary.json"), encoding="utf-8") as f:
            return json.load(f)
    except (OSError, ValueError):
        return None


GAP_SECONDS = 600   # match the iOS conversationGap default: a quiet stretch longer than this is a divider


def _sidecar(rec_dir, fname):
    """The per-clip organized sidecar dict (transcript + started_at + duration_sec), or {}."""
    if not fname:
        return {}
    try:
        with open(os.path.join(rec_dir, ".organized", fname + ".json"), encoding="utf-8") as f:
            return json.load(f)
    except (OSError, ValueError):
        return {}


def _clock(iso):
    """iOS-style gutter label 'h:mmam/pm' from an ISO start, or '' if unparseable."""
    try:
        t = datetime.datetime.fromisoformat(iso)
    except (TypeError, ValueError):
        return ""
    return "%d:%02d%s" % (t.hour % 12 or 12, t.minute, "am" if t.hour < 12 else "pm")


def _quiet_label(seconds):
    mins = int(seconds // 60)
    if mins >= 60:
        return "%d hr %d min quiet" % (mins // 60, mins % 60)
    return "%d min quiet" % mins


def _find_conv(diary, key):
    for day in (diary or {}).get("days", []):
        for conv in day.get("conversations", []):
            clips = conv.get("clips") or []
            if clips and clips[0].get("file", "") == key:
                return conv
    return None


def build(rec_dir, key):
    diary = load_diary(rec_dir)
    conv = _find_conv(diary, key)
    if conv is None:
        return {"ok": False, "error": "Conversation not found."}
    # Blocks = the conversation's transcript, sentence by sentence, across its clips in order.
    # block_meta is index-aligned to blocks (so citation indices stay stable): each clip's FIRST
    # sentence carries the clip's clock time (the gutter timestamp) and, when a long quiet stretch
    # preceded it, a silence-divider label, mirroring the iOS TranscriptFeed manuscript look.
    blocks, block_meta = [], []
    prev_end = None
    for clip in conv.get("clips") or []:
        sc = _sidecar(rec_dir, clip.get("file", ""))
        sents = split_sentences(sc.get("transcript") or "")
        if not sents:
            continue
        started = sc.get("started_at")
        try:
            start_dt = datetime.datetime.fromisoformat(started) if started else None
        except (TypeError, ValueError):
            start_dt = None
        clock = _clock(started) or clip.get("time", "")
        gap = None
        if prev_end is not None and start_dt is not None:
            quiet = (start_dt - prev_end).total_seconds()
            if quiet > GAP_SECONDS:
                gap = _quiet_label(quiet)
        for j, s in enumerate(sents):
            blocks.append(s)
            block_meta.append({"time": clock if j == 0 else "", "gap": gap if j == 0 else None})
        if start_dt is not None:
            prev_end = start_dt + datetime.timedelta(seconds=sc.get("duration_sec") or 0)
    block_tokens = [content_tokens(b) for b in blocks]
    sentences = [{"text": s, "block": best_block(s, block_tokens)}
                 for s in split_sentences(conv.get("summary", ""))]
    return {"ok": True, "title": conv.get("title", ""), "who_with": conv.get("who_with", ""),
            "time_range": conv.get("time_range", ""),
            "summary_sentences": sentences, "blocks": blocks, "block_meta": block_meta}


def main(argv):
    rec_dir = os.path.expanduser(argv[argv.index("--rec-dir") + 1]) if "--rec-dir" in argv else DEFAULT_REC
    key = argv[argv.index("--conv-key") + 1] if "--conv-key" in argv else ""
    sys.stdout.write(json.dumps(build(rec_dir, key), ensure_ascii=False))


if __name__ == "__main__":
    main(sys.argv[1:])
