from ._imports import *
from ._config import config
from ._awsbase import Aws_base

class Aws_comprehend (Aws_base):
    @classmethod
    def _get_client(cls, env):
        return super()._get_client(env, 'comprehend')  

    # Function to replace VINs with {VIN}
    @staticmethod
    def replace_vin(text):
        vin_pattern = r'\b[A-HJ-NPR-Z0-9]{17}\b'
        return re.sub(vin_pattern, 'VIN', text)
    
    
    
    @staticmethod
    def dict_to_dataframe(data_dict, prefix=None):
        """
        Convert a dictionary into a pandas DataFrame where the keys are the column names.
        Optionally, add a prefix to the column names.

        Parameters:
        data_dict (dict): The dictionary to convert.
        prefix (str): The prefix to add to the column names (default is None).

        Returns:
        pd.DataFrame: The resulting DataFrame.
        """
        # Convert the dictionary to a DataFrame
        df = pd.DataFrame([data_dict])
        # Add prefix to column names if provided
        if prefix:
            df = df.add_prefix(prefix)
        return df
    
def client(env):
    return Aws_comprehend._get_client(env)

@atexit.register
def comprehend_close():
    for client in Aws_comprehend.clients.values():
        client.close()


def remove_pii (text, env):
    clean_text = Aws_comprehend.replace_vin(text)
    # reversed to not modify the offsets of other entities when substituting
    response = client(env)[env].detect_pii_entities(
    Text= clean_text,
    LanguageCode='en'
    )
    for NER in reversed(response['Entities']):
        clean_text = clean_text[:NER['BeginOffset']] + NER['Type'] + clean_text[NER['EndOffset']:]
    return clean_text

def get_targeted_sentiment (df, env, join_keys=[]):
    targeted_sentiment = pd.DataFrame()
    error_rows = []

    # Reset the index of the DataFrame in place
    df.reset_index(drop=True, inplace=True)

    for i, row in df.iterrows():    
        txt = row['concatenated_messages'].encode('utf-8')[:5000].decode('utf-8', 'ignore')
        if len(row['concatenated_messages'].encode('utf-8')) > 5000:
            message_length_exceeded = 1
        else:
            message_length_exceeded = 0

        target_sentiment_response = client(env)[env].detect_targeted_sentiment(
            Text=txt,
            LanguageCode='en'
        )

        for ent in target_sentiment_response['Entities']:
            comp_sentiment_score = ent['Mentions'][0]['MentionSentiment']['SentimentScore']
            comp_sentiment_score_df = Aws_comprehend.dict_to_dataframe(comp_sentiment_score, 'comp_sentiment_score_')

            comp_mentions = pd.DataFrame()
            for col in ['Score', 'GroupScore', 'Text', 'Type', 'BeginOffset', 'EndOffset']:
                comp_mentions.loc[0, col] = ent['Mentions'][0][col]
            comp_sentiment = pd.DataFrame({'comp_sentiment': [ent['Mentions'][0]['MentionSentiment']['Sentiment']]})
            temp = comp_mentions.join(comp_sentiment).join(comp_sentiment_score_df)
            for j in join_keys:
                temp[j] = row[j]
            temp['message_length_exceeded'] = message_length_exceeded
            temp = temp[join_keys + ['message_length_exceeded'] + list(comp_mentions.columns) + list(comp_sentiment.columns) + list(comp_sentiment_score_df.columns)]

            targeted_sentiment = pd.concat([targeted_sentiment, temp])        

    return targeted_sentiment, error_rows

