import syntax
import colorsys
from cfg.CFG_Node import CFG_DIAMOND


def _expr_used_names(expr) -> set[str]:
    """Collect variable names (syntax.ID) used inside an expression subtree."""
    used: set[str] = set()

    def visit(node):
        if node is None:
            return
        if isinstance(node, syntax.ID):
            used.add(node.name)
            return
        if isinstance(node, syntax.EXPRESSION):
            for _, child in node.children():
                visit(child)

    visit(expr)
    return used


def _show_analysis_on_node(node) -> bool:
    """Return True if analysis annotations should be displayed for this node."""
    ast = node.ast_node
    if isinstance(node, CFG_DIAMOND):
        return False
    if ast is None:
        return False
    return isinstance(
        ast,
        (
            syntax.ASSIGN,
            syntax.CALL,
            syntax.IF,
            syntax.WHILE,
            syntax.DECL,
            syntax.LET,
            syntax.SEQ,
            syntax.COMP,
            syntax.EQOP,
            syntax.LOP,
        ),
    )


def _lv_in_for_display(node, analysis):
    """Display-level IN set for LV."""
    in_set = set(analysis.in_sets.get(node.id, set()))
    ast_node = node.ast_node
    if isinstance(ast_node, syntax.ASSIGN):
        func = analysis._func_scope.get(node.id)
        rhs_vars = {
            analysis._resolve_var(func, name)
            for name in _expr_used_names(ast_node.expr)
        }
        in_set |= rhs_vars
    return in_set


def _node_color(node_id: int) -> tuple[str, str]:
    """Return (edge_color, fill_color) deterministically for a node id."""
    # Golden-angle hue distribution gives stable, distinct colors.
    hue = ((node_id * 0.6180339887498949) % 1.0)
    edge_rgb = colorsys.hsv_to_rgb(hue, 0.70, 0.82)
    fill_rgb = colorsys.hsv_to_rgb(hue, 0.28, 0.97)

    def to_hex(rgb):
        r, g, b = (int(round(c * 255)) for c in rgb)
        return f"#{r:02x}{g:02x}{b:02x}"

    return to_hex(edge_rgb), to_hex(fill_rgb)


def run_all_analyses(cfg):
    """Run Live Variables and Reached Uses on *cfg*.

    Returns ``(analyses, annotations, ru_edges)`` where:
      • *analyses* is a dict with keys ``"lv"`` and ``"ru"``,
      • *annotations* contains only LivingVariables helper-node labels,
      • *ru_edges* maps definition-node ids to reached use-node ids.
    """
    node_by_id = {n.id: n for n in cfg.nodes()}

    from cfa.live_variables import LiveVariablesAnalysis
    from cfa.reached_uses import ReachedUsesAnalysis

    lv = LiveVariablesAnalysis(cfg)
    ru = ReachedUsesAnalysis(cfg)

    all_ids = set(lv.in_sets.keys()) | set(lv.out_sets.keys())
    annotations = {
        nid: (
            "LivingVariables\\n"
            f"In := {sorted(_lv_in_for_display(node_by_id[nid], lv))}\\n"
            f"Out := {sorted(lv.out_sets.get(nid, set()))}"
        )
        for nid in all_ids
        if lv.in_sets.get(nid, set()) or lv.out_sets.get(nid, set())
        if nid in node_by_id and _show_analysis_on_node(node_by_id[nid])
    }

    return {"lv": lv, "ru": ru}, annotations, ru.reached_uses_by_node()


def analysis_to_dot(
    cfg,
    annotations: dict[int, str],
    analysis_name: str,
    ru_edges: dict[int, list[int]] | None = None,
) -> str:
    """Return a DOT string for *cfg* annotated with analysis results."""
    lines = [
        "digraph CFG {",
        f'  // Analysis: {analysis_name}',
        '  graph [splines=ortho, overlap=false, ranksep=0.7, nodesep=0.45];',
        '  node [fontname="Helvetica"];',
    ]
    color_nodes = set(annotations.keys()) | set((ru_edges or {}).keys())
    node_colors = {nid: _node_color(nid) for nid in color_nodes}

    def emit(node):
        base_label = node.dot_label() or ""
        shape = node.dot_shape
        style = node.dot_style
        style_str = f", {style}" if style else ""
        lines.append(f'  n{node.id} [label="{base_label}", shape={shape}{style_str}];')

        if node.id in annotations:
            ann_id = f"a{node.id}"
            ann_label = annotations[node.id].replace('"', '\\"')
            edge_color, fill_color = node_colors.get(node.id, ("#1f77b4", "#d9ecff"))
            lines.append(
                f'  {ann_id} [label="{ann_label}", shape=note, '
                f'style="filled", fillcolor="{fill_color}", color="{edge_color}", '
                f'fontcolor="{edge_color}"];'
            )
            lines.append(
                f'  {ann_id} -> n{node.id} [style=dotted, arrowhead=none, '
                f'color="{edge_color}"];'
            )

        for i, child in enumerate(sorted(node.children, key=lambda n: n.id)):
            edge_label = ""
            if isinstance(node, CFG_DIAMOND):
                if i == 0:
                    edge_label = ' [label="T"]'
                elif i == 1:
                    edge_label = ' [label="F"]'
            lines.append(f"  n{node.id} -> n{child.id}{edge_label};")

    cfg.traverse(emit, start=cfg.START)

    if ru_edges:
        for idx, def_id in enumerate(sorted(ru_edges)):
            use_ids = sorted(set(ru_edges[def_id]))
            if not use_ids:
                continue

            # One routing hub per definition node to mimic UML-like
            # "out to the side, then down/across to targets" connectors.
            side = "e" if idx % 2 == 0 else "w"
            source_port = "se" if side == "e" else "sw"
            hub_id = f"rh{def_id}"
            edge_color, fill_color = node_colors.get(def_id, ("#1f77b4", "#d9ecff"))

            lines.append(
                f'  {hub_id} [shape=point, width=0.05, height=0.05, '
                f'color="{edge_color}", fillcolor="{edge_color}", style=filled];'
            )
            lines.append(f"  {{ rank=same; n{def_id}; {hub_id}; }}")
            lines.append(
                f'  n{def_id}:{source_port} -> {hub_id} [color="{edge_color}", '
                f'style=dashed, penwidth=1.2, arrowhead=none, constraint=false, '
                f'tailclip=true, headclip=true];'
            )

            for use_id in use_ids:
                if side == "e":
                    target_port = "ne" if (use_id % 2 == 0) else "se"
                else:
                    target_port = "nw" if (use_id % 2 == 0) else "sw"
                lines.append(
                    f'  {hub_id} -> n{use_id}:{target_port} [color="{edge_color}", '
                    f'fontcolor="{edge_color}", fontsize=8, style=dashed, '
                    f'penwidth=1.0, arrowsize=0.6, constraint=false, '
                    f'tailclip=true, headclip=true];'
                )

    lines.append("}")
    return "\n".join(lines)