app_web.py

"""Streamlit web UI for the RAG Research Chatbot."""
import os
import shutil
import zipfile
import streamlit as st
import pandas as pd

from src.config_loader import load_config, get_api_key
from src.ingest import get_chroma_collection, ingest_documents
from src.kb_meta import load_kb_meta_brief
from src.query_engine import understand_query, categorize_query, init_query, _parse_vc_analyzer_result
from src.retriever import clear_collection_cache, retrieve
from src.verifier import verify_and_respond
from src.llm import list_models
from src.stata import explain_
from src.peri import APPROVED_OPTIONS
from src.peri.value_chains import categorize_vc_production, analyze_committment_to_vc , analyze_vc_committment
from src.peri.investments import analyze_investments
from src.prompts import dict_to_string
# from ui import render_db
from huggingface_hub import snapshot_download

snapshot_download(repo_id="CGIAR/peri-kb", 
                  repo_type="dataset", 
                  allow_patterns="ug/*",
                  token=os.getenv('HF_TOKEN'), 
                  local_dir="./"
                  )
shutil.copytree("./ug", "./", dirs_exist_ok=True)
# ── Session initialisation ───────────────────────────────────────────────────

def init_session():
    """Initialise st.session_state with messages list, cfg, and last_retrieval."""
    if "bot_launched" not in st.session_state:
        st.session_state.bot_launched = False
    if "messages" not in st.session_state:
        st.session_state.messages = []
    if "cfg" not in st.session_state:
        try:
            st.session_state.cfg = load_config()
        except Exception as e:
            st.error(f"Configuration error: {e}\n\nPlease run `python setup.py` first.")
            st.stop()
    if "last_retrieval" not in st.session_state:
        st.session_state.last_retrieval = None
    if "pending_clarification" not in st.session_state:
        st.session_state.pending_clarification = None
    if "unresolved_category" not in st.session_state:
        st.session_state.unresolved_category = None
    if "clarification_rounds" not in st.session_state:
        st.session_state.clarification_rounds = 0
    if "resolution_rounds" not in st.session_state:
        st.session_state.resolution_rounds = 0
    if "pending_clarification_question" not in st.session_state:
        st.session_state.pending_clarification_question = None
    if "unresolved_category_question" not in st.session_state:
        st.session_state.unresolved_category_question = None
    if "stata_options" not in st.session_state:
        st.session_state.stata_options = None
    if "stata_file" not in st.session_state:
        st.session_state.stata_file = None

# ── Sidebar ──────────────────────────────────────────────────────────────────
def render_landing_page():
    cfg = st.session_state.cfg
    # print(cfg)
    st.title(f"🤖 Welcome to the {cfg['chatbot'].get('name')} Assistant!")
    st.markdown("Let's go through some preliminary steps below before getting started.")
    
    st.divider()
    
    # Dropdown Options
    country = st.selectbox(
        "Please select a country you would like to prepare a PERI analysis for today:",
        [None]+[c.capitalize() for c in APPROVED_OPTIONS.get("countries")],
        format_func=lambda x: "Please select a country..." if x is None else x
    )

    regime = None
    if country is not None:
        regime = st.selectbox(
            f"Please specify whether {country.capitalize()} is an autocracy or democracy:",
            [None, "Autocracy", "Democracy"],
            format_func=lambda x: f"Please specify {country.capitalize()}'s regime..." if x is None else x
        )
    
    # creativity_level = st.select_slider(
    #     "Adjust Creativity (Temperature):",
    #     options=["Low (Factual)", "Medium (Balanced)", "High (Creative)"],
    #     value="Medium (Balanced)"
    # )
    
    uploaded_ref = None
    if regime is not None:
        # Create the file uploader widget
        pol_geo_desc = """\
            A PERI analysis typically requires a preliminary analysis of 
            the political geography of the country under review. Please upload 
            the political geography analysis results for {country} if available.
            You can also follow the guide at ... to prepare the analysis and upload your results:
            """
        st.write(pol_geo_desc.format(country=country))
        uploaded_ref = st.file_uploader(
            f"Please upload the political geography analysis results for {country}", 
            type=["csv", "xls", "xlsx"])

    if uploaded_ref is not None:
        st.divider()
        
        # Launch Button
        if st.button("Launch PERI-AI 🚀", type="primary", use_container_width=True):
            # Save selections to session state to use during the chat session
            try:
                pol_geo_ref = pd.read_csv(uploaded_ref)
            except Exception:
                pol_geo_ref = pd.read_excel(uploaded_ref)
                
            st.session_state.country = country
            st.session_state.regime = regime
            st.session_state.pol_geo_ref = pol_geo_ref

            with st.status("Preparing metadata for the analysis...", expanded=True) as status:
                status.update(label=f'Preparing value chain metadata...')
                pol_geo_scores = categorize_vc_production(pol_geo_ref, "value_chain")
                
                status.update(label=f'Preparing investment area metadata...')
                investments = analyze_investments()

                c2vc = {}
                for vc in pol_geo_ref["value_chain"].unique():
                    status.update(label=f'Analyzing committment to {vc} production...')
                    prompt = f"Look up and retrieve all relevant information on {vc} production in {country} from the knowledgebase."     
                    qu_result = understand_query(prompt, cfg)

                    search_query = qu_result.get("search_query", prompt)
                    display_query = qu_result.get("display_query", prompt)
                    sql_query = qu_result.get("sql_query")
                    print(sql_query)

                    retrieval_result = retrieve(search_query, cfg, route="both", sql_query=sql_query)
                    st.session_state.last_retrieval = retrieval_result

                    output = verify_and_respond(
                        display_query, retrieval_result, cfg, vc, prompt,
                    )
                    print(output["response"])
                    c2vc.update(_parse_vc_analyzer_result(output["response"]))

                status.update(
                    label=f"Completed metadata preparation for {pol_geo_scores.shape[0]} value chains and {investments.shape[0]} investment areas in {country}.",
                    state="complete",
                    )
            print(c2vc)
            st.session_state.beneficiaries_var = pol_geo_scores
            st.session_state.investments_var = investments
            st.session_state.c2vc = c2vc
            st.session_state.c2vc_results = analyze_vc_committment(c2vc)
            st.session_state.country = country
            st.session_state.bot_launched = True
            st.rerun()



# ── Sidebar ──────────────────────────────────────────────────────────────────

def render_sidebar():
    """Render sidebar with provider/model selectors, web search toggle, and KB stats."""
    cfg = st.session_state.cfg

    with st.sidebar:
        st.header("Settings")

        # --- Provider dropdown ---
        providers = ["openai", "anthropic", "gemini", "meta-llama"]
        current_provider = cfg.get("llm", {}).get("provider", "openai")
        provider_index = providers.index(current_provider) if current_provider in providers else 0

        provider = st.selectbox(
            "LLM Provider",
            providers,
            index=provider_index,
            key="sidebar_provider",
        )

        # Update cfg in session when provider changes
        if provider != cfg.get("llm", {}).get("provider"):
            cfg.setdefault("llm", {})["provider"] = provider

        # --- Model dropdown (cached per provider) ---
        # Invalidate model cache if provider changed
        prev_provider_key = "prev_provider"
        if st.session_state.get(prev_provider_key) != provider:
            for p in providers:
                st.session_state.pop(f"models_{p}", None)
            st.session_state[prev_provider_key] = provider

        models_cache_key = f"models_{provider}"
        if models_cache_key not in st.session_state:
            api_key = get_api_key(cfg, provider)
            if api_key:
                try:
                    st.session_state[models_cache_key] = list_models(provider, api_key)
                except Exception:
                    st.session_state[models_cache_key] = []
            else:
                st.session_state[models_cache_key] = []

        available_models = st.session_state[models_cache_key]
        current_model = cfg.get("llm", {}).get("model", "")

        if available_models:
            model_index = (
                available_models.index(current_model)
                if current_model in available_models
                else 0
            )
            model = st.selectbox(
                "Model",
                available_models,
                index=model_index,
                key="sidebar_model",
            )
        else:
            model = st.text_input(
                "Model",
                value=current_model,
                key="sidebar_model_text",
            )

        # Update cfg in session when model changes
        if model != cfg.get("llm", {}).get("model"):
            cfg.setdefault("llm", {})["model"] = model

        # --- Web search toggle ---
        web_enabled = cfg.get("web_search", {}).get("enabled", False)
        web_toggle = st.toggle("Web search", value=web_enabled, key="sidebar_web_search")
        cfg.setdefault("web_search", {})["enabled"] = web_toggle

        st.divider()

        # --- Knowledge base stats ---
        st.subheader("Knowledge Base")
        try:
            collection = get_chroma_collection(cfg)
            chunk_count = collection.count()
            st.metric("Chunks indexed", chunk_count)
        except Exception as e:
            st.warning(f"Could not read knowledge base: {e}")
            chunk_count = 0

        # --- Re-ingest button ---
        if st.button("Re-ingest documents", use_container_width=True):
            st.info("Ingestion may take a few minutes for large document collections...")
            with st.spinner("Ingesting documents..."):
                try:
                    count = ingest_documents(cfg)
                    clear_collection_cache()
                    st.success(f"Ingested {count} chunks.")
                    # Clear cached data so it refreshes after re-ingest
                    st.session_state.pop("kb_welcome_summary", None)
                    st.rerun()
                except Exception as e:
                    st.error(f"Ingestion failed: {e}")


# ── Chat interface ───────────────────────────────────────────────────────────

def render_chat():
    """Render the chat interface with message history and input."""
    cfg = st.session_state.cfg

    # Display chat history
    for msg in st.session_state.messages:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])

    # Chat input
    user_input = st.chat_input("Ask a question about your knowledge base...",
                            accept_file="multiple", 
                            file_type=["docx", "csv", "xlsx", "xls", "pdf", "rds", "rda", 
                                       "tsv", "sav", "dta", "txt", "md", "json", "do"]
                        )

    if user_input:
        prompt = user_input.text
        uploaded_files = user_input.files
        # Show and store user message
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

        # ── Categorize user query ───────────────────
        save_dir = "uploaded_do_files"
        os.makedirs(save_dir, exist_ok=True)

        # Exclude the just-appended user message to avoid sending
        # the current question twice (once in history, once as query)
        cat_cfg = cfg.get("query_categorization", {})
        max_history = cat_cfg.get("max_history", 6)

        # ── Query understanding ─────────────────────────────────────────
        qu_cfg = cfg.get("query_understanding", {})
        qu_enabled = qu_cfg.get("enabled", True)
        max_history = qu_cfg.get("max_history", 6)

        # ── Check if this is a clarification response ───────────────────
        unresolved = st.session_state.unresolved_category
        if unresolved is not None:
            # This prompt is the user's clarification answer
            # Include the clarification question for context
            unresolved_question = st.session_state.get("unresolved_category_question", "")
            if unresolved_question:
                combined_cat = f"{unresolved} (Clarification: Q: {unresolved_question} A: {prompt})"
            else:
                combined_cat = f"{unresolved} — {prompt}"
            st.session_state.unresolved_category_question = None
            original_query_cat = unresolved
            st.session_state.unresolved_category = None
        else:
            combined_cat = prompt
            original_query_cat = prompt
            st.session_state.resolution_rounds = 0
            unresolved_question = []

        prior_messages = st.session_state.messages[:-1]
        history = [
            {"role": m["role"], "content": m["content"]}
            for m in prior_messages[-max_history:]
        ]
        try:
            qinit_result = init_query(combined_cat, cfg, history)
            print("result: ", qinit_result)
            if qinit_result.get("country", None) is None:
                
                with st.chat_message("assistant"):
                    st.markdown("Please specify the country for which you'd like to run a PERI analysis")
                return
            if qinit_result.get("value_chain", None) is None and qinit_result.get("investment", None) is None:
                
                with st.chat_message("assistant"):
                    st.markdown(f"Please specify the value chain or investment area in {qinit_result.get('country', None)} for which you'd like to run a PERI analysis")
                return
            qcat_result = categorize_query(combined_cat, cfg, history)
            print("result: ", qcat_result)
        except Exception as e:
            print(e)
            qinit_result = {"country": None, "value_chain": None, "investment": None}
            qcat_result = {"category": "pillar_1", "action": "unresolved"}

        esc_char = "\n"
        esc_char1 = "\n -"
        if not isinstance(qinit_result.get("country", None), list) and qinit_result.get("country", None)==None:
            resolution_msg = f"It seems there is no country specified for the analysis. Currently the PERI framework supports analysis for the countries listed below:\n\n  {', '.join([c.capitalize() for c in APPROVED_OPTIONS.get('countries')])}\n\n We are also continuously \
                            working to expand the framework and you can submit a form if the country you would like to run the analysis on is not included. In the meantime please let me know if you would like to run the anlysis for one of the included countries."
                        
        if qinit_result.get("country", None)[0].lower() not in [c.lower() for c in APPROVED_OPTIONS.get('countries')]:
            resolution_msg = f"It seems you are trying to run a PERI analysis for {qinit_result.get('country', None)[0].capitalize()}! Currently the PERI framework only supports analysis for the countries listed below:\n\n  {', '.join([c.capitalize() for c in APPROVED_OPTIONS.get('countries')])}\n\n We are continuously \
                working to expand the framework and you can submit a form to request. In the meantime please let me know if you would like to run the anlysis for one of the included countries."

        if qinit_result.get("value_chain", None)[0]==None and qinit_result.get('investment', None)[0]==None:
            resolution_msg = f"It seems there is no value chain or investment area specified for the analysis. Please select one of the value chains or investment areas included in the current PERI framework."
        
        if qinit_result.get("value_chain", None)[0].lower() not in [c.lower() for v in APPROVED_OPTIONS.get('value_chains').values() for c in v]:
            resolution_msg = f"It seems you are trying to run a PERI analysis for {qinit_result.get('value_chain', None)[0]} in {qinit_result.get('country', None)[0].capitalize()}! Currently the PERI framework only supports analysis for the value chains listed below:\n\n  {dict_to_string(APPROVED_OPTIONS.get('value_chains'), 2)}\n\n We are continuously \
                working to expand the framework and you can submit a form to request. In the meantime please let me know if you would like to run the anlysis for one of the included countries."
                
        if qinit_result.get("investment", None)[0].lower() not in [i.lower() for i in APPROVED_OPTIONS.get('investments')]:
            resolution_msg = f"It seems you are trying to run a PERI analysis for {qinit_result.get('investment', None)[0]} investemnt in {qinit_result.get('country', None)[0].capitalize()}! Currently the PERI framework only supports analysis for the investment areas listed below:\n\n  {', '.join([c.capitalize() for c in APPROVED_OPTIONS.get('investments')])}\n\n We are continuously \
                working to expand the framework and you can submit a form to request. In the meantime please let me know if you would like to run the anlysis for one of the included countries."

        if qinit_result.get("country", None)[0] in APPROVED_OPTIONS.get('countries') and (qinit_result.get('value_chain', None)[0] in APPROVED_OPTIONS.get('value_chains') or qinit_result.get('investment', None)[0] in APPROVED_OPTIONS.get('investments')):
    
            pillar_dict = {
                "pillar_1":f"would like to understand whether {','.join(qinit_result.get('value_chain'))} aligns with the government's political incentives.",
                "pillar_2":f"would like to understand to what degree decisions are impacted by the lobbying of particular groups or by elite influence",
                "pillar_3":f"would like to understand if {','.join(qinit_result.get('value_chain'))} and/or investing in {','.join(qinit_result.get('investment'))} can be feasibly implemented given the broader institutional and policy environment",
                
            }
            resolution_msg = f"**Before I search, could you clarify?** Please let me know if you {' and '.join([pillar_dict[c] for c in qcat_result.get('category')])}."
                                 
        # st.session_state.unresolved_category_question = f"Please let me know if you {' and '.join([pillar_dict[c] for c in qcat_result.get('category')])}."
        st.session_state.unresolved_category_question = resolution_msg
        st.session_state.unresolved_category = original_query_cat
        st.session_state.resolution_rounds += 1
        st.session_state.messages.append({"role": "assistant", "content": resolution_msg})
        
        with st.chat_message("assistant"):
            st.markdown(resolution_msg)
        return

        print("unresolved: ", unresolved)
        print("result: ", qinit_result)
        print("combined: ", combined_cat)
        # if qcat_result.get("category") == "contact_and_info" and qcat_result.get("action") == "contact":
        #     with st.chat_message("assistant"):
        #         st.markdown("For additional assistance, please contact the WEAI Helpdesk at IFPRI-WEAI@cgiar.org.\
        #                      Please let me know if there's anything else I can assist you with today.")
        #     return
        if qcat_result.get("category") == "contact_and_info" and qcat_result.get("action") == "upload":#st.session_state.clarification_rounds < max_clarifications:

            # Allow multiple file uploads
            # uploaded_files = st.file_uploader("Choose files to zip", accept_multiple_files=True)

            if uploaded_files:
                # Specify the path where the zip file will be saved
                zip_filename = "uploaded_files.zip"
                
                # Write uploaded files into a single zip archive on disk
                with zipfile.ZipFile(zip_filename, "w", zipfile.ZIP_DEFLATED) as zipf:
                    for file in uploaded_files:
                        # Write each file's bytes directly into the zip
                        zipf.writestr(file.name, file.getvalue())
                        
                st.success(f"Successfully uploaded {len(uploaded_files)} files!")

            with st.chat_message("assistant"):
                st.markdown("Thank you for sharing these resources with us. The WEAI team will work to \
                            review the files and reach out to you should we need any more information. \
                            Please let me know if there's anything else I can assist you with today.")
            return
            
        if qcat_result.get("category") == "stata" and "unresolved" in qcat_result.get("sub-action") and st.session_state.stata_options is None:
            # Ask clarification — store original query and question for context
            st.session_state.unresolved_category = original_query_cat

            file_path = ""
            if qcat_result.get("action") == "do":
                options = {"rewrite do-file":"rewrite", 
                           "explain the do-file":"explain", 
                           "suggest fix(es) for any errors in the do-file":"suggestfix"}
                # Accept .do files
                uploaded_file = uploaded_files[0]
                if uploaded_file is not None:
                    # Define the full file path on disk
                    file_path = os.path.join(save_dir, uploaded_file.name)

                    # Write the uploaded bytes to the local path
                    with open(file_path, "wb") as f:
                        f.write(uploaded_file.getbuffer())

                    # st.success(f"File successfully uploaded")

            if qcat_result.get("action") == "code":
                options = {"rewrite the stata code":"rewrite", 
                           "explain the stata code":"explain", 
                           "suggest fix(es) for any errors in the code":"suggestfix"}

            if qcat_result.get("action") == "error":
                options = {"explain the error":"explain", 
                           "suggest fix(es) for the error":"suggestfix"}        

            st.session_state.stata_options = options
            st.session_state.stata_file = file_path
            st.session_state.unresolved_category_question = f"Please let me know what you need help with: {' or '.join(options.keys())}. List all options that apply."
            st.session_state.resolution_rounds += 1
            resolution_msg = f"**Before I search, could you clarify?** Please let me know what you need help with: {' or '.join(options.keys())}. List all options that apply."
            st.session_state.messages.append({"role": "assistant", "content": resolution_msg})
            
            with st.chat_message("assistant"):
                st.markdown(resolution_msg)
            return

        print(qcat_result.get("sub-action"))
        if "unresolved" not in qcat_result.get("sub-action", ["unresolved"]):
            st.session_state.unresolved_category = original_query_cat
            st.session_state.resolution_rounds += 1
            st.session_state.messages.append({"role": "assistant", "content": resolution_msg})

            if qcat_result.get("category") == "stata":
                selected_options = [st.session_state.stata_options.get(s) for s in qcat_result.get("sub-action", ["unresolved"])]
                opts = {
                        "rewrite":    True if "rewrite" in selected_options else False,
                        "explain":    True if "explain" in selected_options else False,
                        "suggestfix": True if "suggestfix" in selected_options else False,
                        # "capture":    True if "capture" in selected_options else False,
                        # "verbose":    True if "verbose" in selected_options else False,
                        "lines":      None#lines if "lines" in selected_options else None
                    }
                explanation = explain_(qcat_result.get("action"), combined_cat, cfg, st.session_state.stata_file, opts)
                st.session_state.unresolved_category_question = f"{explanation}\n\n Please let me know if this answer is helpful or if it requires further clarification."
                resolution_msg = f"{explanation}\n\n Please let me know if this answer is helpful or if it requires further clarification."

                with st.chat_message("assistant"):
                    st.markdown(resolution_msg)
                return

        if "unresolved" in qcat_result.get("sub-action", "unresolved") and "Please let me know if this answer is helpful or if it requires further clarification." in unresolved_question:
            st.session_state.pending_clarification = unresolved
            st.session_state.pending_clarification = unresolved_question
            st.session_state.unresolved_category = unresolved
            st.session_state.unresolved_category_question = unresolved_question
            # qu_enabled = False

        # st.markdown(result)
            
        # ── Check if this is a clarification response ───────────────────
        pending = st.session_state.pending_clarification
        if pending is not None:
            # This prompt is the user's clarification answer
            # Include the clarification question for context
            pending_question = st.session_state.get("pending_clarification_question", "")
            if pending_question:
                combined = f"{pending} (Clarification: Q: {pending_question} A: {prompt})"
            else:
                combined = f"{pending} — {prompt}"
            st.session_state.pending_clarification_question = None
            original_query = pending
            st.session_state.pending_clarification = None
        else:
            combined = prompt
            original_query = prompt
            st.session_state.clarification_rounds = 0

        search_query = combined
        display_query = combined
        route = "vector"
        sql_query = None
        max_clarifications = qu_cfg.get("max_clarifications", 1)

        if qu_enabled:
            print("Q understanding")
            # Exclude the just-appended user message to avoid sending
            # the current question twice (once in history, once as query)
            prior_messages = st.session_state.messages[:-1]
            history = [
                {"role": m["role"], "content": m["content"]}
                for m in prior_messages[-max_history:]
            ]
            try:
                qu_result = understand_query(combined, cfg, history)
            except Exception:
                qu_result = {"action": "search", "search_query": combined, "display_query": combined, "original_query": original_query, "route": "vector", "sql_query": None}

            if qu_result.get("action") == "clarify" and st.session_state.clarification_rounds < max_clarifications:
                # Ask clarification — store original query and question for context
                st.session_state.pending_clarification = original_query
                st.session_state.pending_clarification_question = qu_result.get('clarification_question', 'Could you be more specific?')
                st.session_state.clarification_rounds += 1
                clarification_msg = f"**Before I search, could you clarify?** {qu_result.get('clarification_question', 'Could you be more specific?')}"
                st.session_state.messages.append(
                    {"role": "assistant", "content": clarification_msg}
                )
                with st.chat_message("assistant"):
                    st.markdown(clarification_msg)
                return

            # After max clarification rounds, force search (matches CLI behavior)
            if qu_result.get("action") == "clarify":
                qu_result["action"] = "search"

            search_query = qu_result.get("search_query", combined)
            display_query = qu_result.get("display_query", original_query)
            route = qu_result.get("route", "vector")
            sql_query = qu_result.get("sql_query")

        # Generate assistant response
        with st.chat_message("assistant"):
            print("Searching knowledge base")
            with st.status("Searching knowledge base...", expanded=True) as status:
                # Show reformulated query if different
                if search_query != original_query:
                    status.update(label=f'Searching for: "{search_query}"...')

                # Retrieval
                try:
                    retrieval_result = retrieve(search_query, cfg, route=route, sql_query=sql_query)
                except Exception as e:
                    status.update(label=f"Retrieval error: {e}", state="error")
                    st.error(f"Retrieval failed: {e}")
                    return
                st.session_state.last_retrieval = retrieval_result

                n_local = len(retrieval_result.get("db_results", []))
                n_web = len(retrieval_result.get("web_results", []))
                n_sql = len(retrieval_result.get("sql_results", []))
                sql_match = retrieval_result.get("sql_match_type", "")
                source_label = f"Found {n_local} local"
                if n_sql:
                    match_label = f" ({sql_match} match)" if sql_match else ""
                    source_label += f" + {n_sql} SQL rows{match_label}"
                source_label += f" + {n_web} web sources. Generating response..."
                status.update(label=source_label)

                # Response generation uses display_query — a clear,
                # complete question that incorporates any clarification context
                try:
                    result = verify_and_respond(
                        display_query, retrieval_result, cfg,
                        original_query=original_query,
                    )
                except Exception as e:
                    status.update(label=f"Generation error: {e}", state="error")
                    st.error(f"Response generation failed: {e}")
                    return

                # Update status based on verification outcome
                if result.get("refused"):
                    status.update(label="No sufficient sources found.", state="error")
                elif result.get("verification_passed") is True:
                    sql_label = f" + {n_sql} SQL rows" if n_sql else ""
                    status.update(
                        label=f"Verified ({result.get('iterations', 0)} iteration(s)). "
                              f"{n_local} local{sql_label} + {n_web} web sources.",
                        state="complete",
                    )
                elif result.get("verification_passed") is False:
                    status.update(
                        label="Response generated (verification did not fully pass).",
                        state="error",
                    )
                else:
                    sql_label2 = f" + {n_sql} SQL rows" if n_sql else ""
                    status.update(
                        label=f"Done. {n_local} local{sql_label2} + {n_web} web sources.",
                        state="complete",
                    )

            final_response = f"{result.get('response', '')}\n\n For additional assistance, please contact the WEAI Helpdesk at IFPRI-WEAI@cgiar.org.\
                             Please let me know if there's anything else I can assist you with today."
            # Display the response
            st.markdown(final_response)

        # Store assistant message
        st.session_state.messages.append(
            {"role": "assistant", "content": final_response}
        )

        st.session_state.unresolved_category = None
        st.session_state.unresolved_category_question = None
        
        # Cap message history to prevent unbounded memory growth
        max_messages = 200  # 100 Q&A pairs
        if len(st.session_state.messages) > max_messages:
            st.session_state.messages = st.session_state.messages[-max_messages:]

def render_db():
    # rows = st.columns((5,5), gap='medium')
    row0 = st.columns((4,4), gap='medium')
    row1 = st.columns((2,6), gap='medium')
    row2 = st.container()
    row3 = st.container()
    
    with row0[0]:
        st.markdown(f"#### Pillar 1: Value Chain Alignment with Political Incentives")       
        score_cols = ["value_chain", "trend", "strategic_importance_score", 
              "institutional_committment_score", "trend_normalized",
              "strategic_importance_normalized", "institutional_commitment_normalized"]
        
        vc_rename_cols = {
            "value_chain": "Value Chain",
            "noremalized_scores":"Distribution of Beneficiaries",
            "committment_to_vc":"Commitment to Value Chain",
            }
        
        vc_alignment = st.session_state.beneficiaries_var.merge(st.session_state.c2vc_results[score_cols], on="value_chain", how="outer")
        vc_alignment["committment_to_vc"] = vc_alignment[["strategic_importance_normalized", "institutional_commitment_normalized"]].mean(axis=1)
        vc_alignment["Alignment"] = vc_alignment[["noremalized_scores", "committment_to_vc"]].mean(axis=1)
        
        st.dataframe(vc_alignment[["value_chain", "noremalized_scores", "committment_to_vc", "Alignment"]].rename(columns=vc_rename_cols),
                    # column_order=("states", "population"),
                    hide_index=True,
                    width='stretch',
                    # column_config={
                    #     "states": st.column_config.TextColumn(
                    #         "States",
                    #     ),
                    #     "population": st.column_config.ProgressColumn(
                    #         "Population",
                    #         format="%f",
                    #         min_value=0,
                    #         max_value=max(st.session_state.beneficiaries_var.population),
                    #     )}
                    )
    with row0[1]:
        st.markdown(f"#### Pillar 1: Investment Alignment with Political Incentives") 
        investment_rename_cols = {
                "investments":"Investment",
                "time_to_impact":"Time to Impact",
                "targetability":"Targetability",
                "visibility":"Visibility",
                "agg_score":"Investment",
            }
        st.dataframe(st.session_state.investments_var[["investments","time_to_impact", "targetability", "visibility"]].rename(columns=investment_rename_cols),
                    # column_order=("states", "population"),
                    hide_index=True,
                    width='stretch',
                    # column_config={
                    #     "states": st.column_config.TextColumn(
                    #         "States",
                    #     ),
                    #     "population": st.column_config.ProgressColumn(
                    #         "Population",
                    #         format="%f",
                    #         min_value=0,
                    #         max_value=max(st.session_state.beneficiaries_var.population),
                    #     )}
                    )
    with row1[0]:
        st.markdown(f'#### Political Geography of {st.session_state.country}')
        st.dataframe(st.session_state.beneficiaries_var,
                    # column_order=("states", "population"),
                    hide_index=True,
                    width='stretch',
                    # column_config={
                    #     "states": st.column_config.TextColumn(
                    #         "States",
                    #     ),
                    #     "population": st.column_config.ProgressColumn(
                    #         "Population",
                    #         format="%f",
                    #         min_value=0,
                    #         max_value=max(st.session_state.beneficiaries_var.population),
                    #     )}
                    )
            
    with row1[1]:
        st.markdown(f'#### Investment Area Analysis Results')
        st.dataframe(st.session_state.investments_var,
                    # column_order=("states", "population"),
                    hide_index=True,
                    width='stretch',
                    # column_config={
                    #     "states": st.column_config.TextColumn(
                    #         "States",
                    #     ),
                    #     "population": st.column_config.ProgressColumn(
                    #         "Population",
                    #         format="%f",
                    #         min_value=0,
                    #         max_value=max(st.session_state.beneficiaries_var.population),
                    #     )}
                    )  
    with row2:
        st.markdown(f'#### Committment to Value Chain Analysis')
        st.dataframe(st.session_state.c2vc_results,
                    # column_order=("states", "population"),
                    hide_index=True,
                    width='stretch',
                    # column_config={
                    #     "states": st.column_config.TextColumn(
                    #         "States",
                    #     ),
                    #     "population": st.column_config.ProgressColumn(
                    #         "Population",
                    #         format="%f",
                    #         min_value=0,
                    #         max_value=max(st.session_state.beneficiaries_var.population),
                    #     )}
                    )
    with row3:
        # with cols[2]:
        st.markdown(f"#### Committment of {st.session_state.country}'s Government to Various Value Chains")
       
       
        def display_nested_dict(d, indent=0, cols_per_row=2):
        # Get all items in the dictionary
            items = list(d.items())

            # Create rows of columns
            for i in range(0, len(items), cols_per_row):
                cols = st.columns(cols_per_row)
                for j in range(cols_per_row):
                    if i + j < len(items):
                        key, value = items[i + j]
                        with cols[j]:
                            if isinstance(value, dict):
                                with st.expander(f"{' ' * indent}📁 {key.capitalize()}"):
                                # Recursive call (you can adjust columns or keep single column inside expander)
                                    display_nested_dict(value, indent + 2, 1)
                            else:
                                st.write(f"**{key.capitalize()}:** {value}")

        display_nested_dict(st.session_state.c2vc, cols_per_row=3)
        # st.json(st.session_state.c2vc)
# ── Main ─────────────────────────────────────────────────────────────────────

def main():
    """Orchestrate the Streamlit app."""
    st.set_page_config(
        page_title="ResearchBot",
        page_icon="🔬",
        layout="wide",
    )

    init_session()

    cfg = st.session_state.cfg
    bot_name = cfg.get("chatbot", {}).get("name", "ResearchBot")

    st.title(bot_name)

    # Show KB summary on first visit (LLM-generated welcome summary)
    if not st.session_state.messages:
        if "kb_welcome_summary" not in st.session_state:
            st.session_state.kb_welcome_summary = load_kb_meta_brief(cfg)
        kb_summary = st.session_state.kb_welcome_summary
        if kb_summary:
            with st.expander("Knowledge Base Contents", expanded=True):
                st.markdown(kb_summary)
                st.markdown("In case you are familiar with the PERI analysis \
                please let me know what pillar(s) you'd like us to analyze today and for what country. If this is\
                is your first time working with PERI please proceed by letting me know what your particular interests are.")

    if not st.session_state.bot_launched:
        render_landing_page()

    else:
        # render_sidebar()
        # render_chat()
        render_db()

    st.caption(
        "All answers are sourced from the local knowledge base. "
        "Web sources are supplementary only. "
        "Every claim is citation-verified before display."
    )


if __name__ == "__main__":
    main()