import networkx as nx
from fsm_crawl.blocking.labeler import Labeler
from fsm_crawl.utils.domain import fqdn_from_url
from d3graph import d3graph, vec2adjmat

class BrowsingGraph:
    _RED = "#CC1010"
    _GREEN = "#AAFF00"
    
    def __init__(self, labeler: Labeler):
        self.graph = nx.DiGraph()
        self.labeler = labeler

    def add_navigation(self, referrer_url: str | None, target_url: str):
        target_label = self.labeler.label(target_url)
        src_label = self.labeler.label(referrer_url) if referrer_url else None

        dst_fqdn = fqdn_from_url(target_url)
        src_fqdn = fqdn_from_url(referrer_url) if referrer_url else None

        if not dst_fqdn:
            return

        self._add_node(dst_fqdn, target_label)
        if src_fqdn:
            self._add_node(src_fqdn, src_label)
            self.graph.add_edge(src_fqdn, dst_fqdn)
    
    async def add_request(self, request):
        """
        Can be used in async Playwright:
        page.on("request", browsing_graph.add_request)
        """
        url = request.url
        headers = request.headers

        # Use referer or origin as source
        referrer = headers.get("referer") or headers.get("origin") or None
        self.add_navigation(referrer, url)

    def _add_node(self, domain: str, label: str):
        if domain in self.graph:
            return
        color = self._RED if label == 1 else self._GREEN
        self.graph.add_node(domain, color=color)


    def draw_graph(self, path="graph.html"):
        source = []
        target = []
        weight = []

        for u, v, data in self.graph.edges(data=True):
            source.append(u)
            target.append(v)
            weight.append(data.get("weight", 1))

        adjmat = vec2adjmat(source, target, weight=weight)
        d3 = d3graph()
        d3.graph(adjmat)
        colors = [self.graph.nodes[node].get("color", "#AAAAAA") for node in adjmat.columns.values]
        d3.set_node_properties(color=colors, label=adjmat.columns.values)
        d3.show(filepath=f"./dist/{path}")

