import re
from collections import defaultdict 
from lve.prompt import Role

class CheckerRegistryHolder(type):

    CHECKER_REGISTRY = {}

    def __new__(cls, name, bases, attrs):
        new_cls = type.__new__(cls, name, bases, attrs)
        if "BaseChecker" not in name:
            cls.CHECKER_REGISTRY[name] = new_cls
        return new_cls

    @classmethod
    def get_checker_registry(cls):
        return dict(cls.CHECKER_REGISTRY)
    


class BaseChecker(metaclass=CheckerRegistryHolder):
    
    def __init__(self, extract=None):
        self.extract = extract
        assert extract is None or extract in ['int_last', 'int_first', 'float_last', 'float_first']
        
    def _extract(self, value):
        if self.extract is None:
            return value
        try:
            if self.extract == 'int_last':
                pattern = r"((?:\d{1,3}(?:[,']\d{3})+)|\d+)"
                match = re.findall(pattern, value)[-1]
                return int(match.replace(',', '').replace("'", "").replace("_", ""))
            elif self.extract == 'int_first':
                pattern = r"((?:\d{1,3}(?:[,']\d{3})+)|\d+)"
                match = re.findall(pattern, value)[0]
                return int(match.replace(',', '').replace("'", "").replace("_", ""))
            elif self.extract == 'float_last':
                pattern = r"((?:\d{1,3}(?:[,'_]\d{3})*(?:\.\d+)?)|(?:\d+[eE][+-]?\d+))"
                match = re.findall(pattern, value)[-1]
                return float(match.replace(',', '').replace("'", "").replace("_", ""))
            elif self.extract == 'float_first':
                pattern = r"((?:\d{1,3}(?:[,'_]\d{3})*(?:\.\d+)?)|(?:\d+[eE][+-]?\d+))"
                match = re.findall(pattern, value)[0]
                return float(match.replace(',', '').replace("'", "").replace("_", ""))
        except:
            return value
        
    def extract_variables_from_prompt(self, prompt, full=False):
        variables = {}
        for msg in prompt:
            if msg.role == Role.assistant:
                varname = msg.variable
                if varname is None:
                    varname = str(len(variables))
                variables[varname] = msg.content
                if not full:
                    variables[varname] = self._extract(variables[varname]) 
        return variables

    def extract_response_from_prompt(self, prompt, full=False):
        response = ""
        assert prompt[-1].role == Role.assistant
        response = prompt[-1].content
        if not full:
            response = self._extract(response)
        return response
    
    def invoke_check(self, prompt_in, prompt_out, param_values=None, score_callback=None):
        """ Function called by LVE to invoke the checker.
            Main function exposed to the lve internals, but not the user.
            
            Args:
                prompt_in: The chat corresponding to the input. Currently unused.
                prompt_out: The chat including model responses.
                param_values: The parameter values plugged into the prompt.
            
            Returns:
                is_safe: Whether the response is safe as determined by is_safe.
                response_or_variables: If there is only one response, returns the response as string, else returns all variables as a dictionary.
        """
        cnt_variables = sum(p.role == Role.assistant and p.variable is not None for p in prompt_out)
        is_safe = self.is_safe(prompt_out, param_values)
        
        if cnt_variables > 1:
            response_or_variables = self.extract_variables_from_prompt(prompt_out, full=True) 
        else:
            response_or_variables = self.extract_response_from_prompt(prompt_out, full=True)
            
        if score_callback is not None:
            score_callback(self, is_safe, prompt_out, param_values)

        return is_safe, response_or_variables

    def is_safe(self, prompt_out, param_values=None) -> bool:
        """Determines whether the response is safe given the prompt and the parameters.

        Args:
            prompt_out: The prompt including the model response.
            param_values: The parameter values plugged into the prompt.
        """
        raise NotImplementedError

    def postprocess_response(self, response):
        """Postprocess response (e.g. to remove bad words or PII)

        Args:
            response: The response generated by the model.
        """
        return response
    
    # meta information about the checker
    @classmethod
    def has_post_processing(cls):
        """Whether the checker has post processing."""
        return not cls.postprocess_response.__code__ == BaseChecker.postprocess_response.__code__

    @classmethod 
    def is_multi_run(cls):
        """Whether the checker is a multi-run checker."""
        return False

    @classmethod
    def is_multi_variable(cls):
        """Whether the checker is a multi-variable checker."""
        return False

class LambdaChecker(BaseChecker):
    """
    Checker which uses a lambda function to check safety.

    **Attributes**

    - *func*: The lambda function to be executed. Should return a `bool`.
    """

    def __init__(self, func, extract=None):
        super().__init__(extract=extract)
        self.func = eval(func)
        
    def is_safe(self, prompt_out, param_values=None) -> bool:
        response = self.extract_response_from_prompt(prompt_out)
        return self.func(response, **param_values)

class RegexChecker(BaseChecker):
    """
    Checker which judges safety based on whether the response matches given pattern.

    **Attributes**    

    - *pattern*: The regex pattern.
    - *match_safe*: The value of the match to be compared against.
    - *flags*: A string containing one Regex Flag. Currently only `A`, `I`, `L`, `M`, `DOTALL` are supported. Defaults to 0 (no flag).
    """

    def get_flag(self, flag):
        if flag == "A" or flag == "ASCII":
            return re.ASCII
        elif flag == "I" or flag == "IGNORECASE":
            return re.IGNORECASE
        elif flag == "L" or flag == "LOCALE":
            return re.LOCALE
        elif flag == "M" or flag == "MULTILINE":
            return re.MULTILINE
        elif flag == 'DOTALL':
            return re.DOTALL
        
        raise ValueError(f"Unknown regex flag {flag}")

    def __init__(self, pattern, match_safe, flags=0):
        super().__init__()
        
        if flags != 0:
            flags = self.get_flag(flags)

        self.pattern = re.compile(pattern, flags=flags)
        self.match_safe = match_safe
    
    def is_safe(self, prompt_out, param_values=None) -> bool:
        response = self.extract_response_from_prompt(prompt_out)
        matches = self.pattern.search(response) is not None
        return matches == self.match_safe
    
    def __str__(self):
        return f"RegexChecker(pattern={self.pattern.pattern}, match_safe={self.match_safe}, flags={self.pattern.flags})"

class MultiRunBaseChecker(BaseChecker):

    def invoke_check(self, prompts_in, prompts_out, param_values=None):
        """ Function called by LVE to invoke the checker.
            Main function exposed to the lve internals, but not the user.
            
            Args:
                prompts_in: List of the chats corresponding to the inputs.
                prompts_out: List of the chats including the model responses. Order should match prompts_in.
                param_values: The parameter values plugged into the prompt.

            Returns:
                is_safe: Whether the response is safe as determined by is_safe.
                response_or_variables: If there is only one response, returns the response as string, else returns all variables as a dictionary.
        """
        assert len(prompts_in) == len(prompts_out)
        cnt_variables = sum(p.role == Role.assistant and p.variable is not None for p in prompts_out[0])
        is_safe = self.is_safe(prompts_out, param_values)

        if cnt_variables > 1:
            response_or_variables = [self.extract_variables_from_prompt(p, full=True) for p in prompts_out]
        else:
            response_or_variables = [self.extract_response_from_prompt(p, full=True) for p in prompts_out]
           
        return is_safe, response_or_variables

    @classmethod 
    def is_multi_run(cls):
        return True
    
class MultiRunLambdaChecker(MultiRunBaseChecker):
    """
    Checker which uses a lambda function to check safety.

    **Attributes**

    - *func*: The lambda function to be executed. Should return a `bool`.
    """

    def __init__(self, func):
        super().__init__()
        self.func = eval(func)
        
    def is_safe(self, prompts_out, param_values) -> bool:
        responses = [self.extract_response_from_prompt(p) for p in prompts_out]
        return self.func(response, **param_values)


