from tqdm import tqdm
from .prefix_computers import TritonPrefixComputer, CPUPrefixComputer
from .prompt_grouper import PromptGrouper
from .prompt_optimizer import PromptOptimizer


class PromptProcessor:
    """Main class that orchestrates prompt processing and reordering."""
    
    def __init__(self, use_gpu=True, use_parallel_cpu=False):
        if use_gpu:
            try:
                self.prefix_computer = TritonPrefixComputer(fallback_to_cpu=True)
            except:
                self.prefix_computer = CPUPrefixComputer(use_parallel=use_parallel_cpu)
        else:
            self.prefix_computer = CPUPrefixComputer(use_parallel=use_parallel_cpu)
        
        self.grouper = PromptGrouper(self.prefix_computer)
        self.optimizer = PromptOptimizer()
    
    def process_and_reorder(self, reordered_prompts, original_prompts, show_progress=True):
        """
        Main method to process and reorder prompts.
        
        Args:
            reordered_prompts: List of prompts to process
            original_prompts: Original prompts corresponding to reordered_prompts
            show_progress: Whether to show progress bar
        
        Returns:
            tuple: (organized_reordered_prompts, organized_original_prompts, final_index_mapping)
        """
        if not reordered_prompts:
            return [], [], []
        
        if len(reordered_prompts) != len(original_prompts):
            raise ValueError("reordered_prompts and original_prompts must have the same length")
        
        # Group prompts by prefix similarity
        groups = self.grouper.group_by_prefix_similarity(reordered_prompts)
        
        # Optimize each group and calculate scores
        group_scores = []
        iterator = tqdm(groups, desc="Processing groups") if show_progress else groups
        
        for group_indices in iterator:
            optimized_group = self.optimizer.optimize_group_order(reordered_prompts, group_indices)
            score = self._calculate_group_score(reordered_prompts, optimized_group)
            group_scores.append((score, optimized_group))
        
        # Sort groups by score (descending)
        group_scores.sort(key=lambda x: -x[0])
        
        # Create final index mapping
        final_index_mapping = []
        for score, group_indices in group_scores:
            final_index_mapping.extend(group_indices)
        
        # Organize results
        organized_reordered_prompts = [reordered_prompts[idx] for idx in final_index_mapping]
        organized_original_prompts = [original_prompts[idx] for idx in final_index_mapping]
        
        return organized_reordered_prompts, organized_original_prompts, final_index_mapping, group_scores
    
    def _calculate_group_score(self, prompts, group_indices):
        """
        Calculate a score for a group. The score is the length of the prefix 
        common to ALL members multiplied by the number of members in the group.
        This prioritizes larger, more cohesive groups.
        """
        group_size = len(group_indices)
        if group_size <= 1:
            return 0  # Single-item groups have no shared prefix value

        # Get the actual prompts for the current group
        group_prompts = [prompts[idx] for idx in group_indices]

        # Find the length of the prefix common to all prompts in the group
        if not group_prompts or not group_prompts[0]:
            common_prefix_len = 0
        else:
            # The common prefix cannot be longer than the shortest prompt in the group
            min_len = min(len(p) for p in group_prompts)
            common_prefix_len = 0
            for i in range(min_len):
                # Get the item from the first prompt to check against the others
                item_to_check = group_prompts[0][i]
                if all(p[i] == item_to_check for p in group_prompts[1:]):
                    common_prefix_len += 1
                else:
                    # Mismatch found, the common prefix ends here
                    break
        
        # The score is the common prefix length times the group size.
        # This rewards both prefix length and the number of items sharing it.
        score = common_prefix_len * group_size
        return score
    
    @staticmethod
    def _longest_common_prefix_length(list1, list2):
        """Calculate the length of longest common prefix between two lists."""
        length = 0
        for a, b in zip(list1, list2):
            if a == b:
                length += 1
            else:
                break
        return length
    
    def set_computation_method(self, use_gpu=True, use_parallel_cpu=False):
        """Change the computation method."""
        if use_gpu:
            self.prefix_computer = TritonPrefixComputer(fallback_to_cpu=True)
        else:
            self.prefix_computer = CPUPrefixComputer(use_parallel=use_parallel_cpu)
        self.grouper = PromptGrouper(self.prefix_computer)