import pandas as pd
import streamlit as st
from streamlit.logger import get_logger
import langchain


from app_config import ENVIRON
from utils.memory_utils import change_memories
from models.model_seeds import seeds

langchain.verbose = ENVIRON =="dev"
logger = get_logger(__name__)

# TODO: Include more variable and representative names
DEFAULT_NAMES = ["Olivia", "Kit", "Abby", "Tom", "Carolyne", "Jessiny"]

def get_random_name(gender="Neutral", ethnical_group="Neutral", names_df=None):
        if names_df is None:
            names_df = pd.DataFrame(DEFAULT_NAMES, columns=['name'])
            names_df["gender"] = "Neutral"
            names_df["ethnical_group"] = "Neutral"
        
        dfi = names_df
        
        if gender != "Neutral":
            dfi = dfi.query(f"gender=='{gender}'")
        if ethnical_group != "Neutral":
            dfi = dfi.query(f"ethnical_group=='{ethnical_group}'")
        if len(dfi) <=0 :
            dfi = names_df
        return dfi.sample(1)['name'].values[0]

def divide_messages(str_memory, str_ai_prefix="texter", str_human_prefix="helper", include_colon=True):
    message_delimiter = "$%$"
    # Split str memory in messaages according to previous prefix and flatten list
    colon = ":" if include_colon else ""
    str_memory =  f"{message_delimiter}{str_ai_prefix}{colon}".join(str_memory.split(f"{str_ai_prefix}{colon}"))
    str_memory =  f"{message_delimiter}{str_human_prefix}{colon}".join(str_memory.split(f"{str_human_prefix}{colon}"))
    return str_memory.split(message_delimiter)

def add_initial_message(issue, language, memory, str_ai_prefix="texter", str_human_prefix="helper", include_colon=True,
                        texter_name="", counselor_name=""):
    initial_mem_str = seeds.get(issue, "GCT")['memory'].format(counselor_name=counselor_name, texter_name=texter_name)
    message_list = divide_messages(initial_mem_str, str_ai_prefix, str_human_prefix, include_colon)
    colon = ":" if include_colon else ""
    for i, message in enumerate(message_list):
        message = message.strip("\n")
        message = message.strip()
        if message is None or message == "":
            pass
        elif message.startswith(str_human_prefix):
            memory.chat_memory.add_user_message(message.lstrip(f"{str_human_prefix}{colon}").strip())
        elif message.startswith(str_ai_prefix):
            memory.chat_memory.add_ai_message(message.lstrip(f"{str_ai_prefix}{colon}").strip()) 

def create_memory_add_initial_message(memories, issue, language, changed_source=False, texter_name="", counselor_name=""):
    change_memories(memories, language, changed_source=changed_source)

    for memory, _ in memories.items():
        if len(st.session_state[memory].buffer_as_messages) < 1:
            add_initial_message(issue, language, st.session_state[memory], texter_name=texter_name, counselor_name=counselor_name)