#!/usr/bin/env python
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

from collections import defaultdict

# Simplified version of the body info.


class Body(dict):
    def __init__(self, body):
        self["BlockIdKind"] = body["BlockId"]["Kind"]
        if "Variable" in body["BlockId"]:
            self["BlockName"] = body["BlockId"]["Variable"]["Name"][0].split("$")[-1]
        loc = body["Location"]
        self["LineRange"] = (loc[0]["Line"], loc[1]["Line"])
        self["Filename"] = loc[0]["CacheString"]
        self["Edges"] = body.get("PEdge", [])
        self["Points"] = {
            i: p["Location"]["Line"] for i, p in enumerate(body["PPoint"], 1)
        }
        self["Index"] = body["Index"]
        self["Variables"] = {
            x["Variable"]["Name"][0].split("$")[-1]: x["Type"]
            for x in body["DefineVariable"]
        }

        # Indexes
        self["Line2Points"] = defaultdict(list)
        for point, line in self["Points"].items():
            self["Line2Points"][line].append(point)
        self["SrcPoint2Edges"] = defaultdict(list)
        for edge in self["Edges"]:
            src, dst = edge["Index"]
            self["SrcPoint2Edges"][src].append(edge)
        self["Line2Edges"] = defaultdict(list)
        for src, edges in self["SrcPoint2Edges"].items():
            line = self["Points"][src]
            self["Line2Edges"][line].extend(edges)

    def edges_from_line(self, line):
        return self["Line2Edges"][line]

    def edge_from_line(self, line):
        edges = self.edges_from_line(line)
        assert len(edges) == 1
        return edges[0]

    def edges_from_point(self, point):
        return self["SrcPoint2Edges"][point]

    def edge_from_point(self, point):
        edges = self.edges_from_point(point)
        assert len(edges) == 1
        return edges[0]

    def assignment_point(self, varname):
        for edge in self["Edges"]:
            if edge["Kind"] != "Assign":
                continue
            dst = edge["Exp"][0]
            if dst["Kind"] != "Var":
                continue
            if dst["Variable"]["Name"][0] == varname:
                return edge["Index"][0]
        raise Exception("assignment to variable %s not found" % varname)

    def assignment_line(self, varname):
        return self["Points"][self.assignment_point(varname)]