class DocReorderer:
    """Handles prompt reordering based on clustering results."""
    
    def reorder_prompts(self, original_prompts, unique_nodes):
        """Reorder prompts based on clustering tree structure."""
        reordered_prompts = []
        
        for i, original_prompt in enumerate(original_prompts):
            reordered_prompt = self._reorder_single_prompt(i, original_prompt, unique_nodes)
            reordered_prompts.append(reordered_prompt)
        
        return reordered_prompts
    
    def _reorder_single_prompt(self, prompt_index, original_prompt, unique_nodes):
        """Reorder a single prompt based on its position in the tree."""
        original_set = set(original_prompt)
        
        # Find the leaf node containing this prompt
        leaf_node = self._find_leaf_node(prompt_index, unique_nodes)
        if not leaf_node:
            return list(original_prompt)
        
        # If leaf is root or has frequency > 1, use it directly
        if leaf_node.is_root:
            return sorted(list(leaf_node.content))
        
        if leaf_node.frequency > 1:
            prefix_content = leaf_node.content
            prefix_list = sorted(list(prefix_content))
            remaining_elements = original_set - prefix_content
            remaining_list = sorted(list(remaining_elements))
            return prefix_list + remaining_list
        
        # Find best ancestor with frequency > 1
        best_node = self._find_best_ancestor(leaf_node, unique_nodes)
        
        if best_node:
            prefix_content = best_node.content
            prefix_list = sorted(list(prefix_content))
            remaining_elements = original_set - prefix_content
            remaining_list = sorted(list(remaining_elements))
            return prefix_list + remaining_list
        else:
            return list(original_prompt)
    
    def _find_leaf_node(self, prompt_index, unique_nodes):
        """Find the leaf node that contains the given prompt index."""
        for node in unique_nodes.values():
            if prompt_index in node.original_indices and node.is_leaf:
                return node
        return None
    
    def _find_best_ancestor(self, start_node, unique_nodes):
        """
        Find the best ancestor by traversing to the root of the branch,
        returning the highest-level ancestor with a frequency > 1 and non-empty content.
        """
        current_node = start_node
        best_ancestor = None

        while current_node.parent is not None:
            parent_id = current_node.parent
            
            # Handle cases where the parent might not exist in the final tree
            if parent_id not in unique_nodes:
                break
                
            parent_node = unique_nodes[parent_id]
            
            # A valid ancestor must represent a cluster (freq > 1) and have shared content.
            # We keep updating `best_ancestor` as we ascend the tree.
            if parent_node.frequency > 1 and not parent_node.is_empty:
                best_ancestor = parent_node
            
            current_node = parent_node
        
        return best_ancestor