#!/usr/bin/env python3
"""
fb_warc_render.py (v5) — Static renderer for a Facebook WARC/WACZ capture, no
replay engine. Reads captured response bodies directly.

v5 fixes comment threading:
  Facebook loads a post and its comments in SEPARATE network responses, so they
  can't be linked within one response. v5 gathers all responses first, indexes
  every comment by the Facebook ID of the thing it's attached to, then links
  comments to their parent post globally and builds a real
  post -> comment -> reply tree. Comments no longer masquerade as posts.
  Comments that can't be matched to a captured post are listed separately so
  nothing is lost.

Also: reaction/comment/share counts, byte-range video reassembly (+ video_urls.txt),
images, Open Graph fallback for HTML permalink pages.

Usage:  python fb_warc_render.py CAPTURE.wacz [out_dir]
Deps:   warcio, beautifulsoup4   (ffmpeg on PATH for video remux)
"""

import sys, os, json, zipfile, tempfile, html, hashlib, re, glob, subprocess, shutil
from collections import deque, defaultdict
from datetime import datetime, timezone
from warcio.archiveiterator import ArchiveIterator
from bs4 import BeautifulSoup

TIME_KEYS = ("creation_time", "created_time", "publish_time", "created_at", "timestamp")
REACT_KEYS = ("reaction_count", "i18n_reaction_count", "like_count", "reactioncount")
COMMENT_KEYS = ("comment_count", "total_comment_count", "comment_count_reduced")
SHARE_KEYS = ("share_count", "i18n_share_count", "reshares")
COMMENT_CONTAINERS = ("comment_list", "comments", "comment_rendering_instance", "display_comments")
IMG_EXT = (".jpg", ".jpeg", ".png", ".gif", ".webp")


def sha(b):
    return hashlib.sha1(b).hexdigest()[:16]


def find_warcs(path, workdir):
    if path.lower().endswith(".wacz"):
        with zipfile.ZipFile(path) as z:
            z.extractall(workdir)
        return [w for w in glob.glob(os.path.join(workdir, "**", "*.warc*"), recursive=True)
                if w.endswith((".warc", ".warc.gz"))]
    return [path]


def decode(rec):
    try:
        return rec.content_stream().read()
    except Exception:
        return b""


def looks_like_json(txt):
    t = re.sub(r"^for\s*\(;;\);", "", txt.lstrip()[:64]).lstrip()
    return t[:1] in "{["


def parse_fb_json(raw):
    txt = re.sub(r"^for\s*\(;;\);", "", raw.decode("utf-8", "replace").strip()).strip()
    objs = []
    try:
        objs.append(json.loads(txt)); return objs
    except Exception:
        pass
    for line in txt.splitlines():
        line = line.strip()
        if line:
            try: objs.append(json.loads(line))
            except Exception: pass
    return objs


def bfs(node, fn, budget=4000):
    q = deque([node]); seen = 0
    while q and seen < budget:
        cur = q.popleft(); seen += 1
        r = fn(cur)
        if r is not None:
            return r
        if isinstance(cur, dict):
            q.extend(cur.values())
        elif isinstance(cur, list):
            q.extend(cur)
    return None


def nearest_time(node):
    def f(c):
        if isinstance(c, dict):
            for k in TIME_KEYS:
                if isinstance(c.get(k), (int, float)):
                    return c[k]
    return bfs(node, f)


def find_author(node):
    def f(c):
        if isinstance(c, dict):
            for ak in ("actors", "author", "owner"):
                v = c.get(ak)
                if isinstance(v, list) and v and isinstance(v[0], dict) and isinstance(v[0].get("name"), str):
                    return v[0]["name"]
                if isinstance(v, dict) and isinstance(v.get("name"), str):
                    return v["name"]
    return bfs(node, f)


def find_count(node, keys):
    def f(c):
        if isinstance(c, dict):
            for k in c:
                if any(kk in k for kk in keys):
                    v = c[k]
                    if isinstance(v, int):
                        return v
                    if isinstance(v, dict):
                        for ck in ("count", "total_count", "total"):
                            if isinstance(v.get(ck), int):
                                return v[ck]
    return bfs(node, f)


def node_text(node):
    if isinstance(node, dict):
        for container in ("message", "body"):
            c = node.get(container)
            if isinstance(c, dict) and isinstance(c.get("text"), str) and c["text"].strip():
                return c["text"].strip()
    return None


def get_id(node):
    if isinstance(node, dict) and isinstance(node.get("id"), str):
        return node["id"]
    return None


def feedback_ids(node):
    """All ids that could link comments to this post: the post's own id plus any
    'feedback' object ids in its subtree."""
    ids = set()
    nid = get_id(node)
    if nid:
        ids.add(nid)

    def f(c):
        if isinstance(c, dict):
            fb = c.get("feedback")
            if isinstance(fb, dict) and get_id(fb):
                ids.add(fb["id"])
        return None
    bfs(node, f)
    return ids


URL_KEYS = ("wwwURL", "permalink_url", "story_permalink_url", "url", "share_url")
URL_HINTS = ("/posts/", "story_fbid", "pfbid", "/photo", "/videos/", "/permalink", "/reel/")


def find_post_url(node):
    """Find the most permalink-looking facebook.com URL in the post's subtree."""
    best = [None]

    def f(c):
        if isinstance(c, dict):
            for k in URL_KEYS:
                v = c.get(k)
                if isinstance(v, str) and "facebook.com" in v:
                    if any(h in v for h in URL_HINTS):
                        return v                       # strong match, stop
                    if best[0] is None:
                        best[0] = v                    # weak fallback, keep looking
        return None
    hit = bfs(node, f)
    return hit or best[0]


# --- indexing comments globally -------------------------------------------
def index_comments(obj, parent_comments):
    """Record every comment keyed by the id of the container it lives under
    (a feedback id for top-level comments, a comment id for replies)."""
    def rec(node, ancestors=(), depth=0):
        if depth > 90:
            return
        if isinstance(node, dict):
            for ck in COMMENT_CONTAINERS:
                cont = node.get(ck)
                if isinstance(cont, dict) and isinstance(cont.get("edges"), list):
                    pid = None
                    for a in (node,) + tuple(reversed(ancestors[-6:])):
                        pid = get_id(a)
                        if pid:
                            break
                    for e in cont["edges"]:
                        n = e.get("node") if isinstance(e, dict) else None
                        if isinstance(n, dict):
                            t = node_text(n)
                            if t:
                                parent_comments[pid].append(
                                    {"id": get_id(n), "text": t,
                                     "who": find_author(n), "when": nearest_time(n)})
            for v in node.values():
                rec(v, ancestors + (node,), depth + 1)
        elif isinstance(node, list):
            for v in node:
                rec(v, ancestors, depth + 1)
    rec(obj)


def build_comment_tree(parent_id, parent_comments, used, seen):
    out = []
    for c in parent_comments.get(parent_id, []):
        cid = c.get("id")
        if cid and cid in seen:
            continue
        if cid:
            seen.add(cid)
        used.add(parent_id)
        node = dict(c)
        node["replies"] = build_comment_tree(cid, parent_comments, used, seen) if cid else []
        out.append(node)
    return out


def build_posts(all_objs):
    parent_comments = defaultdict(list)
    for obj in all_objs:
        index_comments(obj, parent_comments)

    # set of ids that are comments (so we never emit them as posts)
    comment_ids = {c["id"] for cl in parent_comments.values() for c in cl if c.get("id")}

    posts = []
    seen_text = set()
    used_parents = set()
    seen_comment = set()

    def rec(node, ancestors=(), depth=0):
        if depth > 90:
            return
        if isinstance(node, dict):
            t = node_text(node)
            if t:
                nid = get_id(node)
                dedupe = t[:200]
                if not (nid and nid in comment_ids) and dedupe not in seen_text:
                    seen_text.add(dedupe)
                    match_ids = feedback_ids(node)
                    for a in tuple(reversed(ancestors[-4:])):
                        match_ids |= feedback_ids(a)
                    comments = []
                    for mid in match_ids:
                        comments += build_comment_tree(mid, parent_comments, used_parents, seen_comment)
                    when = who = react = comm = share = None
                    for ctx in (node,) + tuple(reversed(ancestors[-5:])):
                        when = when if when is not None else nearest_time(ctx)
                        who = who if who is not None else find_author(ctx)
                        react = react if react is not None else find_count(ctx, REACT_KEYS)
                        comm = comm if comm is not None else find_count(ctx, COMMENT_KEYS)
                        share = share if share is not None else find_count(ctx, SHARE_KEYS)
                        if when and who and react is not None:
                            break
                    posts.append({"text": t, "when": when, "who": who,
                                  "react": react or 0,
                                  "comm": len(comments) if comments else (comm or 0),
                                  "share": share or 0, "comments": comments,
                                  "url": find_post_url(node)})
            for v in node.values():
                rec(v, ancestors + (node,), depth + 1)
        elif isinstance(node, list):
            for v in node:
                rec(v, ancestors, depth + 1)

    for obj in all_objs:
        rec(obj)

    # comments whose parent post wasn't captured
    leftover = []
    for pid, cl in parent_comments.items():
        if pid not in used_parents:
            for c in cl:
                if c.get("id") not in seen_comment:
                    leftover.append(c)
    return posts, leftover


def og_post(soup, url):
    def meta(p):
        t = soup.find("meta", property=p) or soup.find("meta", attrs={"name": p})
        return t.get("content") if t and t.get("content") else None
    desc, title, img = meta("og:description"), meta("og:title"), meta("og:image")
    if desc or title:
        return {"who": title, "when": None, "text": desc or title, "img": img, "url": url,
                "react": 0, "comm": 0, "share": 0, "comments": []}
    return None


def scripts_json(soup):
    for s in soup.find_all("script"):
        t = (s.get("type") or "").lower()
        if ("json" in t or s.get("data-sjs") is not None) and s.string:
            for o in parse_fb_json(s.string.encode("utf-8", "replace")):
                yield o


def fmt_time(when):
    if not when:
        return "unknown date"
    try:
        if when > 1e12:
            when /= 1000
        return datetime.fromtimestamp(when, tz=timezone.utc).strftime("%Y-%m-%d %H:%M UTC")
    except Exception:
        return str(when)


def reassemble_videos(frag_groups, out, tmp):
    made = []
    for i, (base, frags) in enumerate(frag_groups.items()):
        frags.sort(key=lambda x: x[0])
        raw = b"".join(b for _, b in frags)
        if len(raw) < 1024:
            continue
        rawpath = os.path.join(tmp, f"stream_{i}.mp4")
        with open(rawpath, "wb") as fh:
            fh.write(raw)
        outname = f"video_{i}.mp4"
        outpath = os.path.join(out, "assets", outname)
        r = subprocess.run(["ffmpeg", "-y", "-v", "error", "-i", rawpath, "-c", "copy", outpath],
                           capture_output=True)
        if r.returncode == 0 and os.path.exists(outpath) and os.path.getsize(outpath) > 1024:
            made.append(outname)
        else:
            shutil.copy(rawpath, outpath.replace(".mp4", "_raw.mp4"))
            made.append(outname.replace(".mp4", "_raw.mp4"))
    return made


def strip_range(url):
    return re.sub(r"[?&]byte(start|end)=\d+", "", url)


def render_comment(c, depth=0):
    cw = html.escape(c.get("who") or "Unknown")
    cwhen = html.escape(fmt_time(c.get("when")))
    ct = html.escape(c.get("text", ""))
    inner = "".join(render_comment(r, depth + 1) for r in c.get("replies", []))
    cls = "comment reply" if depth else "comment"
    return (f'<div class="{cls}"><span class=cwho>{cw}</span> '
            f'<span class=cmeta>&middot; {cwhen}</span><br>{ct}{inner}</div>')


HEAD = """<!doctype html><meta charset=utf-8><title>{title}</title>
<style>
 body{{font:16px/1.5 system-ui,sans-serif;max-width:760px;margin:2rem auto;padding:0 1rem;background:#f0f2f5;color:#1c1e21}}
 a{{color:#1877f2}} .post{{background:#fff;border:1px solid #ccd0d5;border-radius:8px;padding:1rem;margin:1rem 0}}
 .meta{{color:#65676b;font-size:.85rem;margin-bottom:.5rem}} .who{{font-weight:600}}
 .text{{white-space:pre-wrap}} img,video{{max-width:100%;border-radius:6px;margin-top:.5rem}}
 .counts{{color:#65676b;font-size:.85rem;margin-top:.5rem;border-top:1px solid #eee;padding-top:.4rem}}
 .fblink{{font-size:.85rem;margin-top:.4rem}}
 .comments{{margin-top:.6rem;border-top:1px solid #eee;padding-top:.4rem}}
 .clabel{{font-size:.8rem;font-weight:600;color:#65676b;margin-bottom:.3rem}}
 .comment{{background:#f0f2f5;border-radius:8px;padding:.5rem .7rem;margin:.4rem 0}}
 .reply{{margin-left:1.4rem;background:#e7e9ed}}
 .cwho{{font-weight:600;font-size:.9rem}} .cmeta{{color:#65676b;font-size:.75rem}}
 .gallery img,.gallery video{{max-width:160px;margin:3px}} h1{{font-size:1.4rem}}
 .summary{{background:#fff;border-radius:8px;padding:1rem;border:1px solid #ccd0d5}}
 .people a{{display:block;padding:.4rem 0}}
</style>
"""


def render_post(p):
    who = html.escape(p.get("who") or "Unknown")
    when = html.escape(fmt_time(p.get("when")))
    text = html.escape(p.get("text", ""))
    img = f'<img src="{html.escape(p["img"])}" loading=lazy>' if p.get("img") else ""
    fblink = (f'<div class=fblink><a href="{html.escape(p["url"])}" target=_blank>View on Facebook ↗</a></div>'
              if p.get("url") else "")
    counts = (f'<div class=counts>👍 {p.get("react",0)} reactions &middot; '
              f'💬 {p.get("comm",0)} comments &middot; ↗ {p.get("share",0)} shares</div>')
    cms = p.get("comments", [])
    chtml = ""
    if cms:
        chtml = (f'<div class=comments><div class=clabel>Comments ({len(cms)})</div>'
                 + "".join(render_comment(c) for c in cms) + "</div>")
    return (f'<div class=post><div class=meta><span class=who>{who}</span> &middot; {when}</div>'
            f'<div class=text>{text}</div>{img}{counts}{fblink}{chtml}</div>')


def slugify(name):
    s = re.sub(r"[^a-z0-9]+", "_", (name or "unknown").lower()).strip("_")
    return s or "unknown"


def main():
    if len(sys.argv) < 2:
        print(__doc__); sys.exit(1)
    src = sys.argv[1]
    out = sys.argv[2] if len(sys.argv) > 2 else "fb_render_out"
    os.makedirs(os.path.join(out, "assets"), exist_ok=True)
    tmp = tempfile.mkdtemp()

    warcs = find_warcs(src, tempfile.mkdtemp())
    if not warcs:
        print("No WARC files found."); sys.exit(1)

    all_objs, images = [], []
    og_posts = []
    vid_frags = defaultdict(list)
    vid_urls = []
    ctypes = {}
    seen = set()

    for warc in warcs:
        with open(warc, "rb") as fh:
            for rec in ArchiveIterator(fh):
                if rec.rec_type != "response":
                    continue
                url = rec.rec_headers.get_header("WARC-Target-URI") or ""
                ct = (rec.http_headers.get_header("Content-Type") or "").lower() if rec.http_headers else ""
                key = ct.split(";")[0] or "unknown"
                ctypes[key] = ctypes.get(key, 0) + 1
                body = decode(rec)
                if not body:
                    continue

                if key.startswith("image/") and "keyframes" not in key:
                    d = sha(body)
                    if d in seen:
                        continue
                    seen.add(d)
                    ext = next((e for e in IMG_EXT if url.lower().split("?")[0].endswith(e)), ".jpg")
                    with open(os.path.join(out, "assets", d + ext), "wb") as im:
                        im.write(body)
                    images.append(d + ext)
                    continue
                if key.startswith("video/") or key.startswith("audio/"):
                    vid_urls.append(url)
                    m = re.search(r"bytestart=(\d+)", url)
                    start = int(m.group(1)) if m else len(vid_frags[strip_range(url)])
                    vid_frags[strip_range(url)].append((start, body))
                    continue

                txt = body.decode("utf-8", "replace")
                if looks_like_json(txt):
                    all_objs.extend(parse_fb_json(body))
                elif "html" in key:
                    soup = BeautifulSoup(body, "html.parser")
                    og = og_post(soup, url)
                    if og:
                        og_posts.append(og)
                    all_objs.extend(scripts_json(soup))
                elif "json" in key or "javascript" in key:
                    all_objs.extend(parse_fb_json(body))

    posts, leftover = build_posts(all_objs)
    posts = og_posts + posts
    videos = reassemble_videos(vid_frags, out, tmp)
    if vid_urls:
        with open(os.path.join(out, "video_urls.txt"), "w") as f:
            f.write("\n".join(sorted(set(vid_urls))))

    # group posts by author -> one page per person
    groups = defaultdict(list)
    for p in posts:
        groups[p.get("who") or "Unknown"].append(p)

    # de-dup author slugs (two different names could collide)
    slugs = {}
    used_slugs = set()
    for author in groups:
        base = slugify(author)
        s = base
        i = 2
        while s in used_slugs:
            s = f"{base}_{i}"; i += 1
        used_slugs.add(s)
        slugs[author] = s

    # one page per person
    for author, plist in groups.items():
        plist.sort(key=lambda p: p.get("when") or 0)
        page = [HEAD.format(title=html.escape(author)),
                f'<p><a href="index.html">&larr; All people</a></p>',
                f'<h1>{html.escape(author)}</h1>',
                f'<div class=summary>{len(plist)} posts archived</div>']
        page += [render_post(p) for p in plist]
        with open(os.path.join(out, f"person_{slugs[author]}.html"), "w", encoding="utf-8") as f:
            f.write("\n".join(page))

    # index / landing page
    P = [HEAD.format(title="Facebook archive"),
         "<h1>Facebook archive</h1>",
         f"""<div class=summary><b>{len(posts)}</b> posts by <b>{len(groups)}</b> people &middot;
 <b>{len(images)}</b> images &middot; <b>{len(videos)}</b> videos &middot;
 <b>{len(leftover)}</b> unlinked comments<br>
 <small>{html.escape(', '.join(f'{k}×{v}' for k,v in sorted(ctypes.items(),key=lambda x:-x[1])))}</small></div>"""]

    P.append("<h1>People</h1><div class='post people'>")
    for author in sorted(groups, key=lambda a: -len(groups[a])):
        P.append(f'<a href="person_{slugs[author]}.html">{html.escape(author)} '
                 f'<small>({len(groups[author])} posts)</small></a>')
    P.append("</div>")

    if leftover:
        P.append("<h1>Comments not linked to a captured post</h1><div class=post>")
        P += [render_comment(c) for c in leftover]
        P.append("</div>")
    if videos:
        P.append("<h1>Reassembled videos</h1><div class='post gallery'>")
        P += [f'<video src="assets/{html.escape(v)}" controls preload=none></video>' for v in videos]
        P.append("</div>")
    if images:
        P.append("<h1>All captured images</h1><div class='post gallery'>")
        P += [f'<img src="assets/{html.escape(fn)}" loading=lazy>' for fn in images]
        P.append("</div>")

    with open(os.path.join(out, "index.html"), "w", encoding="utf-8") as f:
        f.write("\n".join(P))

    print(f"\nDone. {len(posts)} posts by {len(groups)} people, {len(images)} images, "
          f"{len(videos)} videos, {len(leftover)} unlinked comments.")
    if vid_urls:
        print(f"Wrote video_urls.txt ({len(set(vid_urls))} unique stream URLs).")
    print(f"Open: {os.path.join(out, 'index.html')}")


if __name__ == "__main__":
    main()
