try:
    from syntheseus.search.visualization import GraphVizNode, mol_to_image, PURCHASABLE, NOT_PURCHASABLE, NO_PURCHASABLE_INFO
    VISUALIZATION_AVAILABLE = True
except ImportError as e:
    print(f"Warning: Visualization dependencies not available: {e}")
    VISUALIZATION_AVAILABLE = False
    # Define dummy functions or classes if needed
    GraphVizNode = None
    mol_to_image = None
from syntheseus.search.graph.and_or import AndOrGraph, ANDOR_NODE, OrNode, AndNode
from rdkit import Chem
from pathlib import Path
import pprint
from typing import Collection, Optional, List, Tuple
from graphviz import Digraph

def _write_graphviz_graph(
    nodes: List[GraphVizNode],
    edges: List[Tuple[str, str]],
    filename: str,
) -> None:
    # Init visualization graph
    if not filename.endswith(".pdf"):
        raise ValueError("Filename must end in .pdf")
    dotfile_name = filename[:-4]
    G = Digraph("G", filename=dotfile_name)  # strip .pdf
    G.format = "pdf"
    G.attr(rankdir="LR")
    # Draw nodes and edges
    for node in nodes:
        G.node(
            node.id,
            label=node.label,
            fontsize=node.fontsize,
            shape=node.shape,
            style=node.style,
            fillcolor=node.fillcolor,
            color=node.border_color,
        )
    for edge in edges:
        G.edge(edge[0], edge[1], label="")

    # Make pdf
    G.render()

    # Remove intermediate dot file
    dotfile = Path(dotfile_name)
    if dotfile.exists():
        dotfile.unlink()


def visualize_andor(
    graph: AndOrGraph,
    filename: str,
    nodes: Optional[Collection[ANDOR_NODE]] = None,
    draw_mols: bool = True,
) -> None:
    """Visualize an AND/OR graph.

    Args:
        graph: The graph to visualize.
        filename: The filename to save the visualization to. Must end in ".pdf"
        nodes: The nodes to include in the visualization. If None, all nodes are included.
        draw_mols: Whether to draw the molecules in the graph.
    """
    if not VISUALIZATION_AVAILABLE:
        raise RuntimeError("Visualization not available on this system")

    # Extract subgraph
    if nodes is None:
        subgraph = graph._graph
    else:
        subgraph = graph._graph.subgraph(nodes)
    del graph

    # Visualize all nodes
    temp_files = []
    graphviz_nodes = []
    for node in subgraph.nodes:
        rows: List[str] = []
        if isinstance(node, OrNode):
            border_color = PURCHASABLE if node.mol.metadata["is_purchasable"] else NOT_PURCHASABLE
            shape = "ellipse"
            if draw_mols and node.mol.smiles != "":
                img_file = _mol_to_image(Chem.MolFromSmiles(node.mol.smiles))
                temp_files.append(img_file)
                rows.append(f'<TD><IMG SRC="{img_file.name}" SCALE="TRUE"/></TD>')
                del img_file
            else:
                rows.append(f"<TD>SMILES: {node.mol.smiles}</TD>")
        elif isinstance(node, AndNode):
            rows.append("<TD>Reaction</TD>")
            border_color = NO_PURCHASABLE_INFO
            shape = "box"
        else:
            raise TypeError(f"Unknown node type {type(node)}")

        # Metadata
        #node_data_str = pprint.pformat(node.data).replace("\n", "<BR/>")
        node_data_str = ""
        rows.append(f"<TD>{node_data_str}</TD>")

        # Make label
        label = '<<TABLE BORDER="0" CELLBORDER="0" CELLSPACING="0">'
        assert len(rows) > 0
        for row in rows:
            label += f"<TR>{row}</TR>"
        label += "</TABLE>>"
        graphviz_nodes.append(
            GraphVizNode(
                id=str(id(node)),
                label=label,
                border_color=border_color,
                shape=shape,
            )
        )

    # Add edges
    edge_list = []
    for node in subgraph.nodes:
        for child in subgraph.successors(node):
            edge_list.append((str(id(node)), str(id(child))))

    # Make the graph
    _write_graphviz_graph(graphviz_nodes, edge_list, filename)