"""
build_graph.py
--------------
Loads synthetic data from a data directory, applies entity resolution,
builds a heterogeneous NetworkX graph, and saves it to graph.gpickle
in the same directory.

Node types : customer, device, address, payment_instrument
Edge types  : uses_device (strong), pays_with (strong),
              ships_to (weak), same_contact (strong)

Each edge carries: relation, strength, first_seen, last_seen, count, order_ids

Usage:
    python graph/build_graph.py [--data-dir PATH]
    DATA_DIR defaults to data/synthetic/
"""

import argparse
import json
import os
import pickle
import re
from collections import defaultdict

import networkx as nx

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def _load(filename, data_dir):
    with open(os.path.join(data_dir, filename), encoding="utf-8") as f:
        return json.load(f)


# ---------------------------------------------------------------------------
# Entity resolution helpers (3.1 – 3.4)
# ---------------------------------------------------------------------------

def normalize_phone(raw: str) -> str:
    """Strip non-digits; drop a leading country code (91 for India, 1 for US)."""
    digits = re.sub(r"\D", "", raw)
    # Drop leading 91 (India) or 1 (US) if the result would be 10 digits
    for prefix in ("91", "1"):
        if digits.startswith(prefix) and len(digits) - len(prefix) == 10:
            digits = digits[len(prefix):]
            break
    return digits


def normalize_email(raw: str) -> str:
    """Lowercase and strip whitespace."""
    return raw.lower().strip()


def normalize_address(raw: str) -> str:
    """Lowercase, strip whitespace, collapse multiple spaces."""
    return re.sub(r"\s+", " ", raw.lower().strip())


# ---------------------------------------------------------------------------
# Graph construction
# ---------------------------------------------------------------------------

def build_graph(data_dir: str = None) -> nx.Graph:
    if data_dir is None:
        data_dir = os.path.join(BASE_DIR, "data", "synthetic")
    # Load data
    customers         = _load("customers.json", data_dir)
    devices           = _load("devices.json", data_dir)
    addresses         = _load("addresses.json", data_dir)
    payment_instr     = _load("payment_instruments.json", data_dir)
    orders            = _load("orders.json", data_dir)

    # Apply entity resolution (normalizations run on load — see ground rules 3.4)
    # For the synthetic data these are effectively no-ops since fields are already
    # hashed/normalized, but the hooks are live for real-data use.
    for c in customers:
        # phone_hash and email_hash are already SHA256 digests — we normalize
        # the raw strings before hashing in a real pipeline; here we just ensure
        # the lookup maps are keyed consistently.
        pass  # hash fields are already canonical

    # Build lookup maps
    device_map    = {d["device_id"]: d for d in devices}
    addr_map      = {a["address_id"]: a for a in addresses}
    instr_map     = {p["instrument_id"]: p for p in payment_instr}
    cust_map      = {c["customer_id"]: c for c in customers}

    G = nx.Graph()

    # ---------------------------------------------------------------------------
    # Add all nodes with their type attribute (3.5)
    # ---------------------------------------------------------------------------
    for c in customers:
        G.add_node(c["customer_id"], node_type="customer", **{
            k: v for k, v in c.items() if k != "customer_id"
        })

    for d in devices:
        G.add_node(d["device_id"], node_type="device", **{
            k: v for k, v in d.items() if k != "device_id"
        })

    for a in addresses:
        G.add_node(a["address_id"], node_type="address", **{
            k: v for k, v in a.items() if k != "address_id"
        })

    for p in payment_instr:
        G.add_node(p["instrument_id"], node_type="payment_instrument", **{
            k: v for k, v in p.items() if k != "instrument_id"
        })

    # ---------------------------------------------------------------------------
    # Accumulate edge data from orders before writing to graph
    # edge_accum[(u, v)] = {"relation":..., "strength":..., "timestamps":[...], "order_ids":[...]}
    # ---------------------------------------------------------------------------
    EdgeKey = tuple  # (node_a, node_b) always sorted so undirected
    edge_accum: dict[EdgeKey, dict] = defaultdict(lambda: {
        "relation": None, "strength": None, "timestamps": [], "order_ids": []
    })

    def _add_edge_data(u, v, relation, strength, timestamp, order_id):
        key = (min(u, v), max(u, v))
        ea = edge_accum[key]
        ea["relation"] = relation
        ea["strength"] = strength
        ea["timestamps"].append(timestamp)
        ea["order_ids"].append(order_id)

    for order in orders:
        cid  = order["customer_id"]
        did  = order["device_id"]
        iid  = order["payment_instrument_id"]
        aid  = order["address_id"]
        ts   = order["timestamp"]
        oid  = order["order_id"]

        # customer –uses– device (strong)
        _add_edge_data(cid, did, "uses_device", "strong", ts, oid)

        # customer –pays_with– payment_instrument (strong)
        _add_edge_data(cid, iid, "pays_with", "strong", ts, oid)

        # customer –ships_to– address (weak)
        _add_edge_data(cid, aid, "ships_to", "weak", ts, oid)

    # ---------------------------------------------------------------------------
    # customer –same_contact– customer edges (3.6 last bullet)
    # Two customers share the same email_hash OR phone_hash → strong link
    # ---------------------------------------------------------------------------
    email_to_custs: dict[str, list[str]] = defaultdict(list)
    phone_to_custs: dict[str, list[str]] = defaultdict(list)

    for c in customers:
        email_to_custs[c["email_hash"]].append(c["customer_id"])
        phone_to_custs[c["phone_hash"]].append(c["customer_id"])

    same_contact_pairs: set[tuple] = set()
    for group in list(email_to_custs.values()) + list(phone_to_custs.values()):
        if len(group) > 1:
            for i in range(len(group)):
                for j in range(i + 1, len(group)):
                    pair = (min(group[i], group[j]), max(group[i], group[j]))
                    same_contact_pairs.add(pair)

    # Use a synthetic timestamp for same_contact edges (first order timestamp of either cust)
    cust_first_ts: dict[str, str] = {}
    cust_orders: dict[str, list] = defaultdict(list)
    for o in orders:
        cust_orders[o["customer_id"]].append(o)
    for cid, ords in cust_orders.items():
        sorted_ts = sorted(ords, key=lambda x: x["timestamp"])
        cust_first_ts[cid] = sorted_ts[0]["timestamp"]

    for (ca, cb) in same_contact_pairs:
        ts = min(cust_first_ts.get(ca, "2026-01-01T00:00:00"),
                 cust_first_ts.get(cb, "2026-01-01T00:00:00"))
        key = (ca, cb)
        edge_accum[key]["relation"] = "same_contact"
        edge_accum[key]["strength"] = "strong"
        edge_accum[key]["timestamps"].append(ts)
        edge_accum[key]["order_ids"].append("CONTACT_MATCH")

    # ---------------------------------------------------------------------------
    # Write all accumulated edges to the graph (3.7)
    # ---------------------------------------------------------------------------
    for (u, v), data in edge_accum.items():
        sorted_ts = sorted(data["timestamps"])
        G.add_edge(
            u, v,
            relation=data["relation"],
            strength=data["strength"],
            first_seen=sorted_ts[0] if sorted_ts else None,
            last_seen=sorted_ts[-1] if sorted_ts else None,
            count=len(data["order_ids"]),
            order_ids=data["order_ids"],
        )

    return G


# ---------------------------------------------------------------------------
# Save graph (3.8) — try write_gpickle; fall back to pickle.dump if unavailable
# ---------------------------------------------------------------------------
def save_graph(G: nx.Graph, path: str):
    try:
        nx.write_gpickle(G, path)
        return "nx.write_gpickle"
    except AttributeError:
        # networkx >= 3.0 removed write_gpickle — use pickle directly
        with open(path, "wb") as f:
            pickle.dump(G, f, protocol=pickle.HIGHEST_PROTOCOL)
        return "pickle.dump (nx.write_gpickle unavailable in this nx version)"


def load_graph(path: str) -> nx.Graph:
    try:
        return nx.read_gpickle(path)
    except AttributeError:
        with open(path, "rb") as f:
            return pickle.load(f)


# ---------------------------------------------------------------------------
# Console stats (3.9)
# ---------------------------------------------------------------------------
def print_stats(G: nx.Graph):
    from collections import Counter

    node_type_counts = Counter(
        data.get("node_type", "unknown")
        for _, data in G.nodes(data=True)
    )

    edge_rel_counts: dict[str, dict] = defaultdict(lambda: {"strong": 0, "weak": 0, "total": 0})
    for u, v, data in G.edges(data=True):
        rel = data.get("relation", "unknown")
        strength = data.get("strength", "unknown")
        edge_rel_counts[rel][strength] = edge_rel_counts[rel].get(strength, 0) + 1
        edge_rel_counts[rel]["total"] += 1

    cust_degrees = [
        d for n, d in G.degree()
        if G.nodes[n].get("node_type") == "customer"
    ]

    print("\n========== Graph Statistics ==========")
    print("Nodes by type:")
    for ntype, cnt in sorted(node_type_counts.items()):
        print(f"  {ntype:25s}: {cnt}")
    print(f"  {'TOTAL':25s}: {G.number_of_nodes()}")

    print("\nEdges by relation / strength:")
    for rel, counts in sorted(edge_rel_counts.items()):
        strong = counts.get("strong", 0)
        weak   = counts.get("weak", 0)
        total  = counts["total"]
        print(f"  {rel:25s}: total={total:5d}  strong={strong:5d}  weak={weak:5d}")
    print(f"  {'TOTAL':25s}: {G.number_of_edges()}")

    if cust_degrees:
        print(f"\nCustomer node degree — min={min(cust_degrees)}  "
              f"max={max(cust_degrees)}  "
              f"avg={sum(cust_degrees)/len(cust_degrees):.2f}")
    print("======================================")


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--data-dir",
        default=None,
        help="Path to data directory (containing JSON files). Defaults to data/synthetic/",
    )
    args = parser.parse_args()

    data_dir = args.data_dir or os.path.join(BASE_DIR, "data", "synthetic")
    print(f"Building graph from {data_dir} ...")
    G = build_graph(data_dir)

    out_path = os.path.join(data_dir, "graph.gpickle")
    method = save_graph(G, out_path)
    print(f"Graph saved to {out_path}  (method: {method})")

    # Verify reload
    G2 = load_graph(out_path)
    assert G2.number_of_nodes() == G.number_of_nodes(), "Reload node count mismatch!"
    assert G2.number_of_edges() == G.number_of_edges(), "Reload edge count mismatch!"
    print("Reload verified OK.")

    print_stats(G)