import os
import json
import sys
import re
from promptflow.client import load_flow


class PIIEvaluator:
    def __init__(self, model_config):
        current_dir = os.path.dirname(__file__)
        prompty_path = os.path.join(current_dir, "pii.prompty")
        self._flow = load_flow(source=prompty_path, model={"configuration": model_config})

    def __call__(self, *, response: str, **kwargs):
        llm_response = self._flow(response=response)
        try:
            # Try to parse as JSON first
            parsed_response = json.loads(llm_response)
            return parsed_response
        except Exception:
            # If it's not JSON, try to extract the score from XML-like tags
            if isinstance(llm_response, str):
                # Look for <S2>score</S2> pattern
                score_match = re.search(r'<S2>(\d+)</S2>', llm_response)
                # Look for <S1>explanation</S1> pattern
                explanation_match = re.search(r'<S1>(.*?)</S1>', llm_response, re.DOTALL)
                
                if score_match:
                    score = int(score_match.group(1))
                    result_dict = {
                        'PII': score,
                        'score': score,
                        'raw_response': llm_response
                    }
                    
                    # Add explanation if found
                    if explanation_match:
                        explanation = explanation_match.group(1).strip()
                        result_dict['explanation'] = explanation
                        result_dict['reason'] = explanation
                        result_dict['PII_reason'] = explanation
                    
                    return result_dict
            # Fallback: return the raw response
            return {'raw_response': llm_response}