#!/usr/bin/env python3
"""AUDIT-044: M5.5 typed-vptr indirect-reachability survey of the
audit-009 cluster 0x82285000-0x82294000 in sylpheed.db.

READ-ONLY. Do not write to the DB.

Run from xenia-rs root:
    python3 audit-runs/audit-044-m55-cluster-survey/survey.py
"""

import csv
import os
import duckdb

CLUSTER_LO = 0x82285000
CLUSTER_HI = 0x82294000
DB = "sylpheed.db"
OUT = "audit-runs/audit-044-m55-cluster-survey/query_outputs"

# audit-009 L1 PCs
L1_PCS = [
    ("sub_822919C8", 0x822919C8),
    ("sub_82293448", 0x82293448),
    ("sub_82288028", 0x82288028),
    ("sub_82292D80", 0x82292D80),
    ("sub_822851E0", 0x822851E0),
    ("sub_82286BC8", 0x82286BC8),
]

# audit-033 CTOR-PROBE chain (leaf -> root)
AUDIT033_CHAIN = [
    ("sub_82451E20", 0x82451E20),
    ("sub_82450720", 0x82450720),
    ("sub_82450638", 0x82450638),
    ("sub_821CB968", 0x821CB968),
    ("sub_821CD458", 0x821CD458),
    ("sub_821CBEA8", 0x821CBEA8),
    ("sub_821CECF0", 0x821CECF0),
    ("sub_821C4988", 0x821C4988),
]

EXTRA_REFERENCED = [
    ("sub_8228E138", 0x8228E138),
    ("sub_8228E498", 0x8228E498),
    ("sub_82172BA0", 0x82172BA0),
    ("sub_8228A628", 0x8228A628),
]


def main():
    os.makedirs(OUT, exist_ok=True)
    con = duckdb.connect(DB, read_only=True)

    # preload reach sets so we use Python set ops, not nested SQL
    static_reach = {r[0] for r in
                    con.execute("SELECT addr FROM v_reachability_from_entry").fetchall()}
    ind_reach = {r[0] for r in
                 con.execute("SELECT addr FROM v_indirect_reachability_from_entry").fetchall()}
    print(f"static reach |X|={len(static_reach)}; indirect reach |X|={len(ind_reach)};"
          f" newly via M5.5={len(ind_reach - static_reach)}")

    # ---------- Q1/Q2 cluster reach ----------
    q1 = con.execute(
        "SELECT COUNT(DISTINCT f.address) FROM functions f "
        "JOIN v_reachability_from_entry r ON r.addr=f.address "
        "WHERE f.address>=? AND f.address<?",
        [CLUSTER_LO, CLUSTER_HI]).fetchone()[0]
    q2 = con.execute(
        "SELECT COUNT(DISTINCT f.address) FROM functions f "
        "JOIN v_indirect_reachability_from_entry r ON r.addr=f.address "
        "WHERE f.address>=? AND f.address<?",
        [CLUSTER_LO, CLUSTER_HI]).fetchone()[0]
    qt = con.execute("SELECT COUNT(*) FROM functions WHERE address>=? AND address<?",
                     [CLUSTER_LO, CLUSTER_HI]).fetchone()[0]
    print(f"Q1 static-reach in cluster: {q1}/{qt}")
    print(f"Q2 indirect-reach in cluster: {q2}/{qt}")

    # ---------- Q3 newly reachable (CSV + bucket dist) ----------
    nr = con.execute(
        "SELECT f.address, f.name, f.pdata_validated, f.has_eh "
        "FROM functions f "
        "JOIN v_indirect_reachability_from_entry ir ON ir.addr=f.address "
        "LEFT JOIN v_reachability_from_entry sr ON sr.addr=f.address "
        "WHERE f.address>=? AND f.address<? AND sr.addr IS NULL "
        "ORDER BY f.address",
        [CLUSTER_LO, CLUSTER_HI]).fetchall()
    with open(f"{OUT}/q3_newly_reachable.csv", "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["address", "name", "pdata_validated", "has_eh"])
        for r in nr:
            w.writerow([f"0x{r[0]:08x}", r[1], r[2], r[3]])
    print(f"Q3 newly reachable in cluster: {len(nr)} (CSV written)")

    # ---------- Q4 audit-009 L1 PCs ----------
    print("Q4 audit-009 L1 PCs reachability:")
    for name, pc in L1_PCS:
        s = pc in static_reach
        i = pc in ind_reach
        print(f"  {name:<14} 0x{pc:08x}  static={'Y' if s else 'N'} indirect={'Y' if i else 'N'}")

    # ---------- Q5 dispatch sites with cluster targets ----------
    rows = con.execute("""
        SELECT idc.dispatch_pc, ids.vptr_offset, ids.slot, ids.candidate_count,
               idc.vtable_address, idc.method_address, x.source_func, f.name
        FROM indirect_dispatch_candidates idc
        JOIN indirect_dispatch_sites ids ON ids.dispatch_pc=idc.dispatch_pc
        LEFT JOIN xrefs x ON x.source=idc.dispatch_pc AND x.kind='ind_call'
        LEFT JOIN functions f ON f.address=x.source_func
        WHERE idc.method_address>=? AND idc.method_address<?
        ORDER BY ids.candidate_count, idc.dispatch_pc
    """, [CLUSTER_LO, CLUSTER_HI]).fetchall()
    with open(f"{OUT}/q5_dispatch_sites.csv", "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["dispatch_pc", "vptr_offset", "slot", "candidate_count",
                    "vtable_address", "method_address", "source_func", "source_name",
                    "src_static_reach", "src_indirect_reach"])
        for r in rows:
            sf = r[6]
            w.writerow([f"0x{r[0]:08x}", r[1], r[2], r[3],
                        f"0x{r[4]:08x}" if r[4] is not None else "",
                        f"0x{r[5]:08x}",
                        f"0x{sf:08x}" if sf is not None else "",
                        r[7] or "",
                        "Y" if sf in static_reach else "N",
                        "Y" if sf in ind_reach else "N"])
    print(f"Q5 dispatch-site candidate rows: {len(rows)}")
    print(f"   unique dispatch sites: {len({r[0] for r in rows})}")
    src_static = {r[6] for r in rows if r[6] in static_reach}
    src_ind = {r[6] for r in rows if r[6] in ind_reach}
    print(f"   distinct source_funcs static-reach: {len(src_static)}")
    print(f"   distinct source_funcs indirect-reach: {len(src_ind)}")

    # ---------- Q6 vtables with cluster methods ----------
    rows = con.execute("""
        SELECT v.address, v.length, v.class_name, COUNT(*) AS n_in_cluster,
               MIN(m.function_address), MAX(m.function_address)
        FROM vtables v JOIN methods m ON m.vtable_address=v.address
        WHERE m.function_address>=? AND m.function_address<?
        GROUP BY v.address, v.length, v.class_name
        ORDER BY n_in_cluster DESC, v.address
    """, [CLUSTER_LO, CLUSTER_HI]).fetchall()
    with open(f"{OUT}/q6_cluster_vtables.csv", "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["vtable_address", "length", "class_name",
                    "methods_in_cluster", "first_pc", "last_pc"])
        for r in rows:
            w.writerow([f"0x{r[0]:08x}", r[1], r[2], r[3],
                        f"0x{r[4]:08x}", f"0x{r[5]:08x}"])
    print(f"Q6 vtables with cluster methods: {len(rows)}")

    # ---------- audit-033 chain reachability ----------
    print("Audit-033 chain reachability (leaf->root):")
    for name, pc in AUDIT033_CHAIN + EXTRA_REFERENCED + L1_PCS:
        s = pc in static_reach; i = pc in ind_reach
        print(f"  {name:<14} 0x{pc:08x}  static={'Y' if s else 'N'} indirect={'Y' if i else 'N'}")

    # ---------- vptr writer (constructor) reachability ----------
    print("Cluster-vtable vptr-writers (constructors):")
    ctors = con.execute("""
        SELECT vw.vtable_address, vw.writer_pc, vw.writer_function, f.name
        FROM vptr_writes vw LEFT JOIN functions f ON f.address=vw.writer_function
        WHERE vw.vtable_address IN (?, ?)
        ORDER BY vw.vtable_address, vw.writer_pc
    """, [0x820a9c28, 0x820aa024]).fetchall()
    for vt, wpc, wfn, name in ctors:
        sr = wfn in static_reach; ir = wfn in ind_reach
        print(f"  vt=0x{vt:08x} writer_fn=0x{wfn:08x} ({name}) writer_pc=0x{wpc:08x} "
              f"static={'Y' if sr else 'N'} indirect={'Y' if ir else 'N'}")

    con.close()


if __name__ == "__main__":
    main()