import pandas as pd
import streamlit as st

from src.st_helpers import st_setup
from src.common import img_dir, escape_dollars
from src.architectures import *


def show_architecture(architecture: str) -> None:
    """
    Convenience wrapper for the streamlit rendering of an architecture details and the
    ability to interact with the architecture
    :param architecture: the name of the architecture to output
    """
    arch = Architecture.get_architecture(architecture)

    # Segment into two containers for organisation
    arch_container = st.container()
    chat_container = st.container()

    with arch_container:
        st.divider()
        st.write(f'### {arch.name}')
        st.write('#### Architecture description')
        st.write(arch.description)
        if arch.img is not None:
            img = os.path.join(img_dir, arch.img)
            st.image(img, caption=f'Protypical {arch.name} Under Test', width=1000)
        table_data = []
        for j, s in enumerate(arch.steps, start=1):
            table_data.append(
                [j, s.__class__.__name__, s.description, s.config_description()]
            )
        table_cols = ['Step', 'Name', 'Description', 'Config details']
        st.write('#### Architecture pipeline steps')
        st.table(pd.DataFrame(table_data, columns=table_cols))

    with chat_container:
        st.write(f"### Chat with {arch.name}")
        st.write("Note this is a simple single query through the relevant architecture. This is just a sample so you can interact with it and does not manage a chat session history.")

        chat_col, trace_col, request_col = st.columns([3, 2, 2])

        with chat_col:
            with st.chat_message("assistant"):
                st.write("Chat with me in the box below")
        if prompt := st.chat_input("Ask a question"):
            with chat_col:
                with st.chat_message("user"):
                    st.write(prompt)
                request = ArchitectureRequest(query=prompt)
                trace = arch(request)
                with st.chat_message("assistant"):
                    st.write(escape_dollars(request.response))
                with trace_col:
                    st.write("#### Architecture Trace")
                    st.markdown(trace.as_markdown())
                with request_col:
                    st.write("#### Full Request/Response")
                    st.markdown(request.as_markdown())


if st_setup('LLM Arch'):
    st.write("# LLM Architectures")
    Architecture.load_architectures()

    # Display the available architectures
    arch_count = len(Architecture.architectures)
    if arch_count == 1:
        st.write('### 1 Architecture available')
    else:
        st.write(f'### {arch_count} Architectures available')

    if st.button("Force reload of architecture configs"):
        Architecture.load_architectures(force_reload=True)

    arch_names = [a.name for a in Architecture.architectures]
    selected_arch = st.radio(label="Available architectures", label_visibility="hidden", options=arch_names, index=None)
    if selected_arch is None:
        st.info('Select an architecture from above to see details and interact with it')
    else:
        show_architecture(selected_arch)