import streamlit as st
from typing import List, Dict, Union
from PIL import Image

class ChatMemory:
    def __init__(self, max_context_length: int = 2048):
        self.max_context_length = max_context_length
        if "memory" not in st.session_state:
            st.session_state.memory = {
                'history': [],
                'context': []
            }
    
    def update(self, user_input: Union[str, Image.Image], response: str):
        """Store interaction with automatic context pruning"""
        # Store user input
        if isinstance(user_input, Image.Image):
            st.session_state.memory['history'].append(('user', 'image', user_input))
        else:
            st.session_state.memory['history'].append(('user', 'text', user_input))
        
        # Store assistant response
        st.session_state.memory['history'].append(('assistant', 'text', response))
        
        # Maintain context window
        current_length = sum(len(item[2]) for item in st.session_state.memory['history'] if item[1] == 'text')
        while current_length > self.max_context_length and len(st.session_state.memory['history']) > 2:
            removed = st.session_state.memory['history'].pop(0)
            if removed[1] == 'text':
                current_length -= len(removed[2])
    
    def get_context(self) -> str:
        """Generate conversation context string"""
        return "\n".join(
            f"{role}: {content}" 
            for role, type_, content in st.session_state.memory['history']
            if type_ == 'text'
        )
    
    def clear(self):
        st.session_state.memory = {'history': [], 'context': []}