import streamlit as st
import latex2markdown
from langchain.docstore.document import Document
import chromadb
from chromadb.config import Settings
import load_model
from load_model import load_embedding
from load_vectors import load_from_file, load_and_split, create_and_add, load_from_web
persist_directory = load_model.persist_directory

def format_document(document: Document):
    """TODO: Implement a nice style"""
    return document.dict()

def format_result_set(result):
    st.write(latex2markdown.LaTeX2Markdown(result["result"]).to_markdown())

    agree = st.checkbox('Show source documents')
    source_documents = result["source_documents"]
    if agree:
        st.write('Source Documents:')
        for document in source_documents:
            st.write(format_document(document))

#@st.cache_resource
def get_chroma_client():
    return chromadb.Client(Settings(chroma_db_impl="duckdb+parquet",
                                    persist_directory=persist_directory
                                ))
#@st.cache_data
def retrieve_collections():
    client = get_chroma_client()
    all_collections = client.list_collections()
    collections = tuple( [{'name': collection.name, 'model_name': collection.metadata['model_name'], "metadata": collection.metadata} for collection in all_collections] )
    return collections