from ..base_evaluator import BaseEvaluator, EvaluationMetrics
from ...attacks.base_attack import AttackResult
from ...judges.base_judge import BaseJudge
from ...core.registry import evaluator_registry
from typing import List

@evaluator_registry.register("judge")
class JudgeEvaluator(BaseEvaluator):
    """
    An evaluator that uses a judge to evaluate attack results.
    
    This evaluator delegates the evaluation of attack success to a provided judge,
    which makes the determination based on the content of the model's responses.
    """
    
    def __init__(self, judge: BaseJudge, **kwargs):
        """
        Initialize the JudgeEvaluator.
        
        Args:
            judge: The judge to use for evaluating responses
            **kwargs: Additional keyword arguments
        """
        super().__init__(**kwargs)
        self.judge = judge
        self.verbose = kwargs.get("verbose", False)
    
    def evaluate(self, results: List[AttackResult]) -> EvaluationMetrics:
        """
        Evaluate the success of attacks using the provided judge.
        
        Args:
            results: List of attack results to evaluate
            
        Returns:
            EvaluationMetrics containing attack success rate and other metrics
        """
        if not results:
            return EvaluationMetrics(attack_success_rate=0.0)
        
        total_attacks = len(results)
        successful_attacks = 0
        total_scores = 0.0
        for result in results:
            # Set success to False initially since we're re-evaluating
            
            
            # Skip if no final prompt or output
            if not result.final_prompt or not result.output_text:
                continue
            if(result.success== False):
                # Use the judge to evaluate the response
                score,_ = self.judge.judge_prompt(result.final_prompt, result.output_text)
                success = self.judge.is_success(score)
                # Update the result with the judge's evaluation
                result.success = success
                total_scores += score
            else:
                success=1
                total_scores+=self.judge.success_threshold
            if success:
                successful_attacks += 1
                
            if self.verbose:
                print(f"Attack evaluated as {'successful' if success else 'failed'}")
        
        # Calculate attack success rate
        attack_success_rate = successful_attacks / total_attacks if total_attacks > 0 else 0.0
        attack_average_score = total_scores / total_attacks if total_attacks > 0 else 0.0
        if self.verbose:
            print(f"Average score across all attacks: {attack_average_score:.2f}")
        if self.verbose:
            print(f"Total attacks: {total_attacks}, Successful attacks: {successful_attacks}, ASR: {attack_success_rate:.2%}")
        return EvaluationMetrics(attack_success_rate=attack_success_rate,average_attack_score=attack_average_score)