import re
import io
import networkx as nx
import graphviz
from networkx.drawing.nx_pydot import read_dot
import xml.etree.ElementTree as ET


class XMLFTAParser:
    """
    XML Fault Tree Analysis (FTA) parser. Parses the XML file and extracts nodes and edges.
    The file is generated by the Diagrams.net application.
    
    Args:
        xml_file (str): Path to the XML file
        probabilistic (bool): Whether the fault tree is probabilistic
    """
    def __init__(self, xml_file, probabilistic=False):
        self.xml_file = xml_file
        self.tree = ET.parse(xml_file)
        self.root = self.tree.getroot()
        self.nodes = dict()
        self.edges = list()
        self.fault_trees = list()
        self.fault_tree_number = int()
        self.graph_number = int()
        self.probabilistic = probabilistic
        self.node_position = dict()
        self.node_probabilities = dict()

    def parse(self):
        """
        Parse the XML file and extract nodes, edges and node probabilities (if available)
        """
        for element in self.root.findall('.//mxCell'):
            self.parse_nodes(element)
            self.parse_edges(element)
        
        # Parse the node probabilities if the fault tree is probabilistic and get the node position
        if self.probabilistic:
            self.get_node_positions()
            
            for element in self.root.findall('.//mxCell'):
                self.parse_node_probability(element)    

    def parse_nodes(self, element):
        """
        Parse the nodes in the XML file and extract the node ID, value, and style

        Args:
            element (xml.etree.ElementTree.Element): XML element
        """
        id = element.get('id')
        value = element.get('value')
        style = element.get('style')
        
        if style and ('shape=xor' in style.split(';')):
            self.nodes[id] = 'OR'
        elif style and ('shape=or' in style.split(';')):
            self.nodes[id] = 'AND'
        elif style and ('shape=process' in style.split(';')):
            cleaned_value = self.clean_label(value) if value else ''
            self.nodes[id] = f'Action: {cleaned_value}'
        elif style and ('ellipse' in style.split(';')):
            cleaned_value = self.clean_label(value) if value else ''
            self.nodes[id] = cleaned_value
        elif style and ('whiteSpace=wrap' in style.split(';')):
            if 'p' not in value.lower():
                cleaned_value = self.clean_label(value) if value else ''
                self.nodes[id] = cleaned_value
        else:
            return

    def parse_edges(self, element):
        """
        Parse the edges in the XML file and extract the source and target nodes

        Args:
            element (xml.etree.ElementTree.Element): XML element
        """
        source = element.get('source')
        target = element.get('target')
        
        if source and target:
            self.edges.append((source, target))
            
    def parse_node_probability(self, element):
        """
        Parse the node probability from the label

        Args:  
            element (xml.etree.ElementTree.Element): XML element
        """
        style = element.get('style')
        
        if style and ('text' in style.split(';')):
            probability = element.get('value')
            
            # Remove blank spaces and convert to float
            if "=" in probability:
                _, probability = probability.split("=")
            probability = float(probability.strip())
            
            # Get the closest node ID
            node_id = self.get_closest_node_id(element)
            self.node_probabilities[node_id] = probability
            
    def get_node_positions(self):
        """
        Get the positions of the nodes in the XML file
        """
        for node_id in self.nodes.keys():
            for element in self.root.findall(f'.//mxCell[@id="{node_id}"]/mxGeometry'):
                x = element.get('x')
                y = element.get('y')
                self.node_position[node_id] = (float(x), float(y))
        
    def get_closest_node_id(self, element):
        """
        Get the closest node ID to the given element

        Args:
            element (xml.etree.ElementTree.Element): XML element

        Returns:
            str: Closest node ID
        """
        element = element.find('mxGeometry')
        x = float(element.get('x'))
        y = float(element.get('y'))
        min_distance = float('inf')
        closest_node_id = None
        
        for node_id, (node_x, node_y) in self.node_position.items():
            distance = (x - node_x) ** 2 + (y - node_y) ** 2
            if distance < min_distance:
                min_distance = distance
                closest_node_id = node_id
        
        return closest_node_id
        
    def clean_label(self, label):
        """
        Remove HTML tags and their content from the label

        Args:
            label (str): Label of the node

        Returns:
            str: Cleaned label
        """
        label = re.sub(r'<[^>]*>', '', label)
        
        return label
    
    def clean_graph(self, G):
        """
        Clean the NetworkX graph by removing nodes with no edges

        Args:
            G (networkx.classes.digraph.DiGraph): NetworkX graph

        Returns:
            networkx.classes.digraph.DiGraph: Cleaned NetworkX graph
        """
        G.remove_nodes_from(list(nx.isolates(G)))
        
        return G
    
    def extract_subtree(self, root_node):
        """
        Extract the subtree for a given root node.

        Args:
            root_node: The ID of the root node.

        Returns:
            A tuple containing the nodes and edges of the subtree.
        """
        subtree_nodes = {root_node: self.nodes[root_node]}
        subtree_edges = set()
        queue = [root_node]

        while queue:
            current_node = queue.pop(0)
            for edge in self.edges:
                if edge[1] == current_node:
                    source, _ = edge
                    subtree_edges.add(edge)
                    if source not in subtree_nodes:
                        subtree_nodes[source] = self.nodes[source]
                        queue.append(source)

        return subtree_nodes, subtree_edges

    def create_graphviz_dot(self, root_node):
        """
        Create a Graphviz dot string from the nodes and edges
        
        Args:
            root_node: The ID of the root node.

        Returns:
            str: Graphviz dot string
            networkx.classes.digraph.DiGraph: NetworkX graph
        """
        self.graph_number += 1
        subtree_nodes, subtree_edges = self.extract_subtree(root_node)
        dot_str = f"digraph G{self.graph_number}" + "{\n"
    
        # Add nodes to the dot string, ensuring node IDs are quoted
        for node_id, node_label in subtree_nodes.items():
            quoted_node_id = f'"{node_id}"'
            if node_label == "OR":
                dot_str += f'    {quoted_node_id} [label="OR", shape="triangle"];\n'
            elif node_label == "AND":
                dot_str += f'    {quoted_node_id} [label="AND", shape="invtriangle"];\n'
            elif node_label.startswith("Action"):
                cleaned_label = node_label.replace('"', '\\"').replace("<", "&lt;").replace(">", "&gt;")
                dot_str += f'    {quoted_node_id} [label="{cleaned_label}", shape="ellipse"];\n'
            else:
                # Clean the label by escaping double quotes and replacing < and > with HTML entities
                cleaned_label = node_label.replace('"', '\\"').replace("<", "&lt;").replace(">", "&gt;")
                dot_str += f'    {quoted_node_id} [label="{cleaned_label}", shape="box"];\n'
        
        # Add edges to the dot string
        for source, target in subtree_edges:
            quoted_source = f'"{source}"'
            quoted_target = f'"{target}"'
            dot_str += f'    {quoted_source} -> {quoted_target};\n'
        
        dot_str += "}"
        
        # Read the dot string into a NetworkX graph
        dot_file = io.StringIO(dot_str)
        G = read_dot(dot_file)
        
        # Clean the graph
        G = self.clean_graph(G)
            
        return dot_str, G
            
    def render_dot_string(self, dot_str, view=False):
        """
        Render the Graphviz dot string

        Args:
            dot_str (str): Graphviz dot string
            view (bool): Whether to view the rendered graph
        """
        self.fault_tree_number += 1
        graphviz.Source(dot_str).render(f'output_fault_tree_{self.fault_tree_number}.gv', view=view)
        
    def find_root_nodes(self):
        """
        Find the root nodes in the fault tree. These are the nodes that are not targets of any edges.

        Returns:
            list: A list of root nodes.
        """
        all_targets = {target for _, target in self.edges}
        root_nodes = [node for node in self.nodes if node not in all_targets]
        
        return root_nodes
    
    def find_leaf_nodes(self):
        """
        Find the leaf nodes in the fault tree. These are the nodes that do not appear as a source in any edge.

        Returns:
            list: A list of leaf nodes.
        """
        all_sources = {source for source, _ in self.edges}
        leaf_nodes = [node for node in self.nodes if node not in all_sources]

        return leaf_nodes
    
    def generate_fault_trees(self, plot=False):
        """
        Generate the tree from the XML file. The tree is represented as a list of NetworkX graphs.
        
        Args:
            plot (bool): Whether to plot the tree using NetworkX
        
        Returns:
            list: A list of NetworkX graphs representing the fault trees
        """
        # Parse the XML file
        self.parse()
        root_nodes = self.find_leaf_nodes()
        
        for root in root_nodes:
            dot_str, G = self.create_graphviz_dot(root)
            # Create a graphviz dot render from the dot string if render is True
            if plot:
                self.render_dot_string(dot_str, view=plot)
                
            self.fault_trees.append(G)
            
        return self.fault_trees