"""Shared model constants, hub routing helpers, and heterogeneous GAT blocks."""

import dgl
import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl.nn import GATConv


EDGE_TYPE_MAPPING = {
    "ret": 0,
    "jumptable": 1,
    "indirectcall": 3,
    "tailcall": 2,
}

EDGE_TYPE_TO_FNHASH = {
    "ret": "trainrethash",
    "jumptable": "trainjthash",
    "tailcall": "trainitchash",
    "indirectcall": "trainicallhash",
}

TASK_TO_CANDIDATE_KEY = {
    "indirectcall": "icall",
    "tailcall": "itailcall",
    "ret": "ret",
    "jumptable": "jumptable",
    "icall": "icall",
    "itailcall": "itailcall",
}

BASE_RELATION_TYPES = (
    "code2code_edges",
    "code2funchead",
    "codecall_edges",
    "codejump_edges",
    "codexrefcode_edges",
    "codexrefdata_edges",
    "dataxrefcode_edges",
    "dataxrefdata_edges",
    "code_to_gch",
    "gch_to_code",
    "data_to_gdh",
    "gdh_to_data",
    "gch_to_gdh",
    "gdh_to_gch",
)

PAPER_HUB_RELATION_TYPES = (
    "src_to_gch",
    "gch_to_src",
    "dst_to_gch",
    "gch_to_dst",
    "data_to_gdh",
    "gdh_to_data",
    "gch_to_gdh",
    "gdh_to_gch",
)

PAPER_RELATION_TYPES = tuple(
    dict.fromkeys(
        tuple(rel for rel in BASE_RELATION_TYPES if rel not in {
            "code_to_gch",
            "gch_to_code",
            "data_to_gdh",
            "gdh_to_data",
            "gch_to_gdh",
            "gdh_to_gch",
        })
        + PAPER_HUB_RELATION_TYPES
    )
)

XREF_RELATION_TYPES = {
    "codexrefcode_edges",
    "codexrefdata_edges",
    "dataxrefcode_edges",
    "dataxrefdata_edges",
}

PAPER_HUB_INCOMING_RELATIONS = {
    "gch_to_src",
    "gch_to_dst",
    "gdh_to_data",
}

HUB_KIND_GCH = 0
HUB_KIND_GDH = 1
MODEL_ARTIFACT_FORMAT = "icflownet_model_v1"


def save_model_artifact(path, model, model_config=None):
    """Save weights together with the model configuration needed to load them."""
    model_config = dict(model_config or {})
    for attribute in (
        "code_feat_dim",
        "data_feat_dim",
        "gdh_radius",
        "use_gch",
        "use_gdh",
        "task_aware_routing",
        "use_reverse_edges",
        "hidden_dim",
        "num_heads",
        "num_layers",
        "shared_layers",
        "task_layers",
        "dropout",
        "task_weights",
        "task_type",
    ):
        if hasattr(model, attribute):
            model_config.setdefault(attribute, getattr(model, attribute))
    artifact = {
        "artifact_format": MODEL_ARTIFACT_FORMAT,
        "model_config": model_config,
        "model_state_dict": model.state_dict(),
    }
    decision_threshold = getattr(model, "decision_threshold", None)
    if decision_threshold is not None:
        artifact["decision_threshold"] = float(decision_threshold)
    torch.save(artifact, path)


def load_model_artifact(model, path, map_location=None):
    """Load a versioned model artifact or training checkpoint."""
    artifact = torch.load(path, map_location=map_location)

    metadata = {}
    if isinstance(artifact, dict) and artifact.get("artifact_format") == MODEL_ARTIFACT_FORMAT:
        model_config = dict(artifact.get("model_config", {}))
        # Older model files included this redundant field in model_config.
        model_config.pop("architecture_version", None)
        stored_gdh_radius = model_config.get("gdh_radius")
        requested_gdh_radius = getattr(model, "gdh_radius", None)
        if (
            stored_gdh_radius is not None
            and requested_gdh_radius is not None
            and int(stored_gdh_radius) != int(requested_gdh_radius)
        ):
            raise RuntimeError(
                "Checkpoint GDH radius {} does not match requested radius {}.".format(
                    stored_gdh_radius,
                    requested_gdh_radius,
                )
            )
        state_dict = artifact["model_state_dict"]
        decision_threshold = float(artifact.get("decision_threshold", 0.5))
        if not 0.0 <= decision_threshold <= 1.0:
            raise RuntimeError(
                "Checkpoint decision threshold must be between 0 and 1."
            )
        model.decision_threshold = decision_threshold
        metadata = {
            "model_config": model_config,
            "artifact_format": artifact.get("artifact_format"),
            "decision_threshold": decision_threshold,
        }
    elif isinstance(artifact, dict) and "model_state_dict" in artifact:
        state_dict = artifact["model_state_dict"]
        metadata = {
            "artifact_format": "training_checkpoint",
        }
    else:
        raise RuntimeError(
            "Unsupported unversioned checkpoint. Expected a versioned model artifact "
            "or training checkpoint."
        )

    try:
        model.load_state_dict(state_dict)
    except RuntimeError as exc:
        raise RuntimeError(
            "Failed to load checkpoint weights: {}".format(exc)
        ) from exc
    return metadata


def normalize_task_name(task_name):
    """Map training task names to the candidate sidecar's keys."""
    return TASK_TO_CANDIDATE_KEY.get(task_name, task_name)


def extract_node_features(g):
    """Collect per-node-type input features using the repo's existing conventions."""
    feat_dict = {}

    for ntype in g.ntypes:
        node_data = g.nodes[ntype].data
        if "featmean" in node_data:
            feat_dict[ntype] = node_data["featmean"]
        elif "feat" in node_data:
            feat_dict[ntype] = node_data["feat"]
        elif ntype == "hub" and "hub_kind" in node_data:
            # Paper-exact hubs use nn.Embedding rather than binary-derived input
            # features.  The model consumes hub_kind directly.
            continue
        elif g.num_nodes(ntype) == 0:
            continue
        else:
            raise RuntimeError(
                f"Missing 'featmean' or 'feat' in node features for node type '{ntype}'."
            )

    if "code" not in feat_dict:
        raise RuntimeError("Missing required 'code' node features.")

    return feat_dict


def add_paper_data_structural_features(g, feat_dict):
    """Append normalized base-graph in/out degrees to saved data features.

    Released graphs persist the normalized-address feature only.  The paper also
    describes lightweight structural data features, so paper_exact derives the
    available in/out-degree signals at runtime without changing the saved graph.
    Hub edges are excluded so these features remain properties of the base ACFG.
    """
    if "data" not in feat_dict or "data" not in g.ntypes:
        return feat_dict

    data_features = feat_dict["data"]
    num_data_nodes = g.num_nodes("data")
    in_degree = data_features.new_zeros(num_data_nodes)
    out_degree = data_features.new_zeros(num_data_nodes)

    for canonical_etype in g.canonical_etypes:
        src_type, rel_type, dst_type = canonical_etype
        if rel_type in PAPER_HUB_RELATION_TYPES:
            continue
        if src_type == "data":
            out_degree += g.out_degrees(etype=canonical_etype).to(
                device=data_features.device,
                dtype=data_features.dtype,
            )
        if dst_type == "data":
            in_degree += g.in_degrees(etype=canonical_etype).to(
                device=data_features.device,
                dtype=data_features.dtype,
            )

    def normalize_degree(values):
        values = torch.log1p(values)
        maximum = values.max() if values.numel() else values.new_zeros(())
        return values / maximum.clamp_min(1.0)

    structural = torch.stack(
        (normalize_degree(in_degree), normalize_degree(out_degree)),
        dim=1,
    )
    augmented = dict(feat_dict)
    augmented["data"] = torch.cat((data_features, structural), dim=1)
    return augmented


def _valid_node_ids(base_graph, values, ntype):
    """Return in-range integer node ids without trusting sidecar contents."""
    if ntype not in base_graph.ntypes:
        return set()

    max_nodes = base_graph.num_nodes(ntype)
    valid_ids = set()
    for value in values or ():
        try:
            node_id = int(value)
        except (TypeError, ValueError):
            continue
        if 0 <= node_id < max_nodes:
            valid_ids.add(node_id)
    return valid_ids


def collect_xref_data_neighborhood(base_graph, code_candidates, radius=2):
    """Collect data nodes within an undirected xRef-only hop radius.

    The paper defines GDH pruning by shortest-path distance in the recovered
    xRef subgraph.  Existing artifact generation treated code->data and
    data->code xRefs symmetrically for the one-hop case, so this implementation
    preserves that convention while extending it to a real bounded BFS.
    """
    if radius < 1 or "code" not in base_graph.ntypes or "data" not in base_graph.ntypes:
        return set()

    start_nodes = {
        ("code", node_id)
        for node_id in _valid_node_ids(base_graph, code_candidates, "code")
    }
    if not start_nodes:
        return set()

    adjacency = {}

    def add_neighbor(left, right):
        adjacency.setdefault(left, set()).add(right)

    for canonical_etype in base_graph.canonical_etypes:
        src_type, rel_type, dst_type = canonical_etype
        if rel_type not in XREF_RELATION_TYPES:
            continue

        src_nodes, dst_nodes = base_graph.edges(etype=canonical_etype)
        src_values = src_nodes.detach().cpu().tolist()
        dst_values = dst_nodes.detach().cpu().tolist()
        for src_id, dst_id in zip(src_values, dst_values):
            src_key = (src_type, int(src_id))
            dst_key = (dst_type, int(dst_id))
            add_neighbor(src_key, dst_key)
            add_neighbor(dst_key, src_key)

    visited = set(start_nodes)
    frontier = set(start_nodes)
    data_nodes = set()

    for _depth in range(1, radius + 1):
        next_frontier = set()
        for node_key in frontier:
            for neighbor in adjacency.get(node_key, ()):
                if neighbor in visited:
                    continue
                visited.add(neighbor)
                next_frontier.add(neighbor)
                if neighbor[0] == "data":
                    data_nodes.add(neighbor[1])
        frontier = next_frontier
        if not frontier:
            break

    return data_nodes


def _cached_xref_data_neighborhood(
    base_graph,
    candidate_meta,
    cache_name,
    code_candidates,
    radius,
):
    """Cache topology-only GDH membership inside an in-memory metadata dict."""
    runtime_cache = candidate_meta.setdefault("_paper_exact_runtime_cache", {})
    signature = (
        int(radius),
        int(base_graph.num_nodes("code")) if "code" in base_graph.ntypes else 0,
        int(base_graph.num_nodes("data")) if "data" in base_graph.ntypes else 0,
        tuple(sorted(int(node_id) for node_id in code_candidates)),
    )
    cached = runtime_cache.get(cache_name)
    if cached is not None and cached.get("signature") == signature:
        return set(cached.get("data_candidates", ()))

    data_candidates = collect_xref_data_neighborhood(
        base_graph,
        code_candidates,
        radius=radius,
    )
    runtime_cache[cache_name] = {
        "signature": signature,
        "data_candidates": sorted(data_candidates),
    }
    return data_candidates


def _task_code_candidate_sets(base_graph, candidate_meta, candidate_key):
    code_meta = candidate_meta.get("code_candidates", {})
    src_candidates = _valid_node_ids(
        base_graph,
        code_meta.get("src", {}).get(candidate_key, []),
        "code",
    )
    dst_candidates = _valid_node_ids(
        base_graph,
        code_meta.get("dst", {}).get(candidate_key, []),
        "code",
    )

    if candidate_key == "jumptable":
        dst_by_src = code_meta.get("dst_by_src", {}).get("jumptable", {})
        for dst_nodes in dst_by_src.values():
            dst_candidates.update(_valid_node_ids(base_graph, dst_nodes, "code"))

    return src_candidates, dst_candidates


def _build_paper_augmented_graph(
    base_graph,
    src_candidates,
    dst_candidates,
    data_candidates,
    use_gch,
    use_gdh,
):
    """Materialize the paper's role-separated Dual-Hub topology."""
    use_gch = bool(use_gch and "code" in base_graph.ntypes)
    use_gdh = bool(use_gdh and "data" in base_graph.ntypes)
    if not use_gch and not use_gdh:
        return base_graph
    if "hub" in base_graph.ntypes:
        raise ValueError("Paper hub augmentation expects an unaugmented base graph.")
    if not src_candidates and not dst_candidates and not data_candidates:
        return base_graph

    graph_device = base_graph.device
    graph_dtype = base_graph.idtype
    graph_data = {}
    for canonical_etype in base_graph.canonical_etypes:
        graph_data[canonical_etype] = base_graph.edges(etype=canonical_etype)

    num_nodes_dict = {ntype: base_graph.num_nodes(ntype) for ntype in base_graph.ntypes}
    hub_node_ids = {}
    hub_kinds = []

    if use_gch:
        hub_node_ids["gch"] = len(hub_kinds)
        hub_kinds.append(HUB_KIND_GCH)
    if use_gdh:
        hub_node_ids["gdh"] = len(hub_kinds)
        hub_kinds.append(HUB_KIND_GDH)

    num_nodes_dict["hub"] = len(hub_kinds)

    def ids_tensor(values):
        return torch.tensor(list(values), dtype=graph_dtype, device=graph_device)

    def add_code_hub_relation(nodes, forward_rel, reverse_rel):
        if not nodes:
            return
        sorted_nodes = sorted(int(node_id) for node_id in nodes)
        gch_id = hub_node_ids["gch"]
        graph_data[("code", forward_rel, "hub")] = (
            ids_tensor(sorted_nodes),
            ids_tensor([gch_id] * len(sorted_nodes)),
        )
        graph_data[("hub", reverse_rel, "code")] = (
            ids_tensor([gch_id] * len(sorted_nodes)),
            ids_tensor(sorted_nodes),
        )

    if use_gch:
        add_code_hub_relation(src_candidates, "src_to_gch", "gch_to_src")
        add_code_hub_relation(dst_candidates, "dst_to_gch", "gch_to_dst")

    if use_gdh and data_candidates:
        sorted_data = sorted(int(node_id) for node_id in data_candidates)
        gdh_id = hub_node_ids["gdh"]
        graph_data[("data", "data_to_gdh", "hub")] = (
            ids_tensor(sorted_data),
            ids_tensor([gdh_id] * len(sorted_data)),
        )
        graph_data[("hub", "gdh_to_data", "data")] = (
            ids_tensor([gdh_id] * len(sorted_data)),
            ids_tensor(sorted_data),
        )

    if use_gch and use_gdh:
        graph_data[("hub", "gch_to_gdh", "hub")] = (
            ids_tensor([hub_node_ids["gch"]]),
            ids_tensor([hub_node_ids["gdh"]]),
        )
        graph_data[("hub", "gdh_to_gch", "hub")] = (
            ids_tensor([hub_node_ids["gdh"]]),
            ids_tensor([hub_node_ids["gch"]]),
        )

    augmented_graph = dgl.heterograph(
        graph_data,
        num_nodes_dict=num_nodes_dict,
        idtype=graph_dtype,
        device=graph_device,
    )

    for ntype in base_graph.ntypes:
        for feat_name, feat_tensor in base_graph.nodes[ntype].data.items():
            augmented_graph.nodes[ntype].data[feat_name] = feat_tensor

    augmented_graph.nodes["hub"].data["hub_kind"] = torch.tensor(
        hub_kinds,
        dtype=torch.long,
        device=graph_device,
    )
    return augmented_graph


def build_paper_task_augmented_graph(
    base_graph,
    candidate_meta,
    task_name,
    use_gch=False,
    use_gdh=False,
    task_aware_routing=True,
    gdh_radius=2,
):
    """Build the paper-exact single-task topology from a reusable base graph."""
    if not use_gch and not use_gdh:
        return base_graph
    if not candidate_meta:
        return base_graph

    candidate_key = normalize_task_name(task_name)
    src_candidates, dst_candidates = _task_code_candidate_sets(
        base_graph,
        candidate_meta,
        candidate_key,
    )
    # Jump-table dispatch sites still route to the source side of the code
    # hub.  Only their destination candidates are excluded by the paper's
    # task-aware topology.  Keep this identical to the unified MTL path below.
    if task_aware_routing and candidate_key == "jumptable":
        dst_candidates = set()

    data_candidates = _cached_xref_data_neighborhood(
        base_graph,
        candidate_meta,
        "single:{}".format(candidate_key),
        src_candidates | dst_candidates,
        gdh_radius,
    ) if use_gdh else set()

    return _build_paper_augmented_graph(
        base_graph,
        src_candidates,
        dst_candidates,
        data_candidates,
        use_gch,
        use_gdh,
    )


def build_paper_mtl_augmented_graph(
    base_graph,
    candidate_meta,
    use_gch=False,
    use_gdh=False,
    task_aware_routing=True,
    gdh_radius=2,
):
    """Build one unified paper MTL graph shared by all task branches."""
    if not use_gch and not use_gdh:
        return base_graph
    if not candidate_meta:
        return base_graph

    src_candidates = set()
    dst_candidates = set()
    for candidate_key in ("icall", "itailcall", "ret", "jumptable"):
        task_src, task_dst = _task_code_candidate_sets(
            base_graph,
            candidate_meta,
            candidate_key,
        )
        # Keep jump-table sources: angr cannot reliably distinguish them from
        # indirect-tail-call sources, and both belong to the shared
        # ``jumptable_itailcall`` source group recorded in hub metadata.
        src_candidates.update(task_src)
        if candidate_key != "jumptable" or not task_aware_routing:
            dst_candidates.update(task_dst)

    data_candidates = _cached_xref_data_neighborhood(
        base_graph,
        candidate_meta,
        "mtl",
        src_candidates | dst_candidates,
        gdh_radius,
    ) if use_gdh else set()

    return _build_paper_augmented_graph(
        base_graph,
        src_candidates,
        dst_candidates,
        data_candidates,
        use_gch,
        use_gdh,
    )


class PaperHeteroGATLayer(nn.Module):
    """Paper-exact HGAT layer with local-conditioned hub-gated fusion."""

    def __init__(
        self,
        in_dim,
        out_dim,
        num_heads,
        dropout=0.4,
        use_residual=True,
        use_reverse_edges=False,
    ):
        super(PaperHeteroGATLayer, self).__init__()

        assert out_dim % num_heads == 0, "out_dim must be divisible by num_heads"

        self.in_dim = in_dim
        self.out_dim = out_dim
        self.num_heads = num_heads
        # ``out_dim`` is the total hidden width.  Split it evenly across the
        # attention heads and concatenate their outputs back to ``out_dim``.
        self.head_dim = out_dim // num_heads
        self.use_reverse_edges = use_reverse_edges
        self.use_residual = use_residual

        relation_types = list(PAPER_RELATION_TYPES)
        if use_reverse_edges:
            relation_types.extend(
                "rev_" + rel_type
                for rel_type in PAPER_RELATION_TYPES
                if rel_type not in PAPER_HUB_RELATION_TYPES
            )
        relation_types = tuple(dict.fromkeys(relation_types))

        self.gat_layers = nn.ModuleDict({
            rel_type: GATConv(
                in_dim,
                self.head_dim,
                num_heads,
                dropout,
                dropout,
                allow_zero_in_degree=True,
            )
            for rel_type in relation_types
        })

        if use_residual:
            self.residuals = nn.ModuleDict({
                ntype: nn.Linear(in_dim, out_dim)
                for ntype in ("code", "data")
            })
            # The implementation stores both virtual nodes under DGL's compact
            # ``hub`` node type, but retains the paper's distinct gcode/gdata
            # residual transforms through their hub_kind identities.
            self.hub_residuals = nn.ModuleList([
                nn.Linear(in_dim, out_dim),
                nn.Linear(in_dim, out_dim),
            ])

        ordinary_types = ("code", "data")
        self.hub_gates = nn.ModuleDict({
            ntype: nn.Linear(out_dim + in_dim, 1)
            for ntype in ordinary_types
        })
        self.hub_scales = nn.ParameterDict({
            ntype: nn.Parameter(torch.ones(()))
            for ntype in ordinary_types
        })
        self.layer_norms = nn.ModuleDict({
            ntype: nn.LayerNorm(out_dim)
            for ntype in ordinary_types
        })

    def _zero_message(self, nfeat):
        return nfeat.new_zeros((nfeat.shape[0], self.out_dim))

    def _sum_messages(self, messages, nfeat):
        if not messages:
            return self._zero_message(nfeat)
        return torch.stack(messages, dim=0).sum(dim=0)

    def forward(self, g, feat_dict):
        local_messages = {ntype: [] for ntype in feat_dict}
        hub_messages = {ntype: [] for ntype in feat_dict}

        for canonical_etype in g.canonical_etypes:
            src_type, rel_type, dst_type = canonical_etype
            if src_type not in feat_dict or dst_type not in feat_dict:
                continue
            if rel_type.startswith("rev_") and not self.use_reverse_edges:
                continue
            if rel_type not in self.gat_layers or g.num_edges(canonical_etype) == 0:
                continue

            relation_heads = self.gat_layers[rel_type](
                g[canonical_etype],
                (feat_dict[src_type], feat_dict[dst_type]),
            )
            relation_message = relation_heads.flatten(start_dim=1)

            if dst_type in ("code", "data") and rel_type in PAPER_HUB_INCOMING_RELATIONS:
                hub_messages[dst_type].append(relation_message)
            else:
                local_messages[dst_type].append(relation_message)

        final_features = {}
        for ntype, nfeat in feat_dict.items():
            local_message = self._sum_messages(local_messages[ntype], nfeat)
            if self.use_residual and ntype in self.residuals:
                residual = self.residuals[ntype](nfeat)
            elif self.use_residual and ntype == "hub":
                hub_kinds = g.nodes["hub"].data["hub_kind"].long()
                residual = self._zero_message(nfeat)
                for kind, transform in enumerate(self.hub_residuals):
                    kind_mask = (hub_kinds == kind).to(nfeat.dtype).unsqueeze(-1)
                    residual = residual + kind_mask * transform(nfeat)
            else:
                residual = self._zero_message(nfeat)

            if ntype in ("code", "data"):
                hub_message = self._sum_messages(hub_messages[ntype], nfeat)
                gate = torch.sigmoid(
                    self.hub_gates[ntype](torch.cat([local_message, nfeat], dim=-1))
                )
                combined = (
                    local_message
                    + self.hub_scales[ntype] * gate * hub_message
                    + residual
                )
                final_features[ntype] = F.elu(self.layer_norms[ntype](combined))
            else:
                # Hub nodes use the standard relation-aware update without the
                # ordinary-node local/hub gate from Equations (3)-(4).
                final_features[ntype] = F.elu(local_message + residual)

        return final_features