import streamlit as st
from streamlit.logger import get_logger

from langchain.memory import ConversationBufferMemory
from models.databricks.instruction_memory import CustomBufferInstructionMemory
from utils.mongo_utils import new_convo

logger = get_logger(__name__)

def push_convo2db(memories, username, language):
    if len(memories) == 1:
        issue =  memories['memory']['issue']
        model_one = memories['memory']['source']
        new_convo(st.session_state['db_client'], issue, language, username, False, model_one)
    else:
        issue =  memories['commonMemory']['issue']
        model_one = memories['memoryA']['source']
        model_two = memories['memoryB']['source']
        new_convo(st.session_state['db_client'], issue, language, username, True, model_one, model_two)

def change_memories(memories, language, changed_source=False):
    for memory, params in memories.items():
        if (memory not in st.session_state) or changed_source:
            source = params['source']
            logger.info(f"Source for memory {memory} is {source}")
            if source in ('OA_rolemodel','OA_finetuned',"CTL_llama2"):
                st.session_state[memory] = ConversationBufferMemory(ai_prefix='texter', human_prefix='helper')
            elif source in ('CTL_mistral'):
                st.session_state[memory] = CustomBufferInstructionMemory(human_prefix="</s> [INST]", memory_key="history")
    
    if ("convo_id" in st.session_state) and changed_source:
        del st.session_state['convo_id']
    

def clear_memory(memories, username, language):
    for memory, _ in memories.items():
        st.session_state[memory].clear()

    if "convo_id" in st.session_state:
        del st.session_state['convo_id']