def save_chat_message(client, collection, role, text, user_id=None):
    """
    Save a single chat message to the Chroma collection.
    Metadata: role ('user' or 'agent'), ts (timestamp), source='chat', optional user_id.
    """
    if not collection:
        return
    ts = float(time.time())
    # create an id unique enough: chat_userid_timestamp_ms_rand
    doc_id = f"chat_{user_id or 'anon'}_{int(ts*1000)}"
    try:
        collection.add(
            documents=[text],
            metadatas=[{
                "role": role,
                "ts": ts,
                "source": "chat",
                **({"user_id": user_id} if user_id else {})
            }],
            ids=[doc_id]
        )
        # persist if client supports it
        try:
            if hasattr(client, "persist"):
                client.persist()
        except Exception:
            pass
    except Exception as e:
        # non-fatal: log to console so UI doesn't break
        print("Error saving chat message to Chroma:", e)

def load_user_chat(collection, user_id):
    """
    Load all chat messages for given user_id from Chroma collection.
    Returns a list of (role, text) ordered by ts ascending.
    Looks for metadata 'source' == 'chat' or 'role' in metadata.
    """
    if not collection or not user_id:
        return []

    docs = []
    metas = []

    # Try preferred 'get' API first
    try:
        res = collection.get(where={"user_id": user_id})
        # some Chroma clients return dict with 'documents' and 'metadatas'
        docs = res.get("documents", [])
        metas = res.get("metadatas", [])
    except Exception:
        # fallback: large query filtered by metadata
        try:
            res = collection.query(query_texts=[""], n_results=1000, where={"user_id": user_id})
            docs = res.get("documents", [[]])[0]
            metas = res.get("metadatas", [[]])[0]
        except Exception as e:
            print("Error loading chat (fallback):", e)
            return []

    # sanitize shapes
    if not isinstance(docs, list) or not isinstance(metas, list):
        return []

    combined = []
    for doc, meta in zip(docs, metas):
        if not isinstance(meta, dict):
            continue
        # accept chat items marked with source='chat' OR those having a role field
        if meta.get("source") == "chat" or meta.get("role") in ("user", "agent"):
            try:
                ts = float(meta.get("ts", 0))
            except Exception:
                ts = 0.0
            role = meta.get("role", "agent" if meta.get("source") != "chat" else "user")
            # doc might be a string, but sometimes nested lists; coerce
            if isinstance(doc, list):
                # some query outputs nest results; flatten join
                doc_text = " ".join(doc)
            else:
                doc_text = str(doc)
            combined.append((ts, role, doc_text))

    # sort by timestamp ascending and return (role, text) list
    combined.sort(key=lambda x: x[0])
    history = [(role, text) for (_ts, role, text) in combined]
    return history

def submit():
    user_message = st.session_state['user_input']
    if user_message:
        st.session_state.history.append(("user", user_message))
        save_chat_message(client, collection, role="user", text=user_message, user_id=username) # save user's message
        
        result = run_agent_once(user_message, client, collection, user_id=username)
        st.session_state.history.append(("agent", result["formatted"]))
        
        save_chat_message(client, collection, role="agent", text=result["formatted"], user_id=username) #save agent message
        if 'empty_key' not in st.session_state:
            st.session_state.empty_key=''
        st.session_state.empty_key=st.session_state.user_input
        # if you set user_input to a blank valueto clear the prompt, streamlit throws an error
        st.session_state.user_input = ''

username = (username or "").strip() or "Default user"

if 'active_profile' not in st.session_state or st.session_state.active_profile != username:
    st.session_state.active_profile = username
    # load saved chat for this username and populate session history
    try:
        if username:
            CHROMA_DIR = f".chroma_db_{username}"
        else:
            CHROMA_DIR = ".chroma_db"

        client, collection = init_chromadb(CHROMA_DIR)
        loaded_history = load_user_chat(collection, username)
        # replace in-memory history with loaded one (or empty list if none)
        st.session_state.history = loaded_history if loaded_history else []
    except Exception as e:
        print("Error loading user chat into session_state:", e)
        st.session_state.history = []