import streamlit as st
import pandas as pd
from comparison import compare_ingredient_methods_ui
from ui_core import embeddings, load_examples
from ui_ingredient_matching import categorize_products
from ui_category_matching import categorize_products_by_category
from ui_hybrid_matching import categorize_products_with_voyage_reranking
from ui_expanded_matching import categorize_products_with_openai_reranking
# Removed unused import: from ui_formatters import format_results_html

# Initialize session state keys if they don't exist
if 'ingredient_input' not in st.session_state:
    st.session_state.ingredient_input = ""
if 'category_input' not in st.session_state:
    st.session_state.category_input = ""
if 'voyage_input' not in st.session_state:
    st.session_state.voyage_input = ""
if 'openai_input' not in st.session_state:
    st.session_state.openai_input = ""
if 'compare_input' not in st.session_state:
    st.session_state.compare_input = ""


def render_ui():
    """Render the Streamlit interface"""
    # Page config is now set in app.py
    st.title("Product Categorization Tool")
    st.markdown("Analyze products by matching to ingredients or categories using AI embeddings.")

    # Use st.tabs for the different sections
    tab_ingredient, tab_category, tab_voyage, tab_openai, tab_compare = st.tabs([
        "Ingredient Embeddings",
        "Category Embeddings",
        "Voyage AI Reranking",
        "OpenAI Reranking",
        "Compare Methods"
    ])

    # --- Ingredient Matching Tab ---
    with tab_ingredient:
        st.header("Match Products to Ingredients")
        col1, col2 = st.columns(2)
        with col1:
            # Handle button click *before* rendering the text area
            if st.button("Load Examples", key="ingredient_examples"):
                st.session_state.ingredient_input = load_examples() # Update state for next rerun

            # Input section - Use the session state value
            text_input = st.text_area(
                "Product Names (one per line)",
                value=st.session_state.ingredient_input, # Use value from state
                placeholder="Enter product names, one per line",
                height=250,
                key="ingredient_input_widget" # Use a different key for the widget itself if needed, or manage via value
            )
            # Update session state if user types manually
            st.session_state.ingredient_input = text_input

            use_expansion = st.checkbox(
                "Use Description Expansion (AI)",
                value=False,
                key="ingredient_expansion",
                help="Expand product descriptions using AI before matching"
            )
            top_n = st.slider("Top N Results", 1, 25, 10, step=1, key="ingredient_top_n")
            confidence = st.slider("Similarity Threshold", 0.1, 0.9, 0.5, step=0.05, key="ingredient_confidence")

            find_ingredients_btn = st.button("Find Similar Ingredients", type="primary", key="ingredient_find")

        with col2:
            # Results section
            st.subheader("Results")
            results_placeholder_ingredient = st.empty()
            if find_ingredients_btn:
                if st.session_state.ingredient_input: # Check state value
                    results_html = categorize_products(
                        st.session_state.ingredient_input,
                        False,
                        use_expansion,
                        top_n,
                        confidence
                    )
                    results_placeholder_ingredient.markdown(results_html, unsafe_allow_html=True)
                else:
                    results_placeholder_ingredient.warning("Please enter product names.")

    # --- Category Matching Tab ---
    with tab_category:
        st.header("Match Products to Categories")
        col1, col2 = st.columns(2)
        with col1:
            if st.button("Load Examples", key="category_examples"):
                st.session_state.category_input = load_examples()

            category_text_input = st.text_area(
                "Product Names (one per line)",
                value=st.session_state.category_input,
                placeholder="Enter product names, one per line",
                height=250,
                key="category_input_widget"
            )
            st.session_state.category_input = category_text_input

            category_use_expansion = st.checkbox(
                "Use Description Expansion (AI)",
                value=False,
                key="category_expansion",
                help="Expand product descriptions using AI before matching"
            )
            category_top_n = st.slider("Top N Categories", 1, 10, 5, step=1, key="category_top_n")
            category_confidence = st.slider("Matching Threshold", 0.1, 0.9, 0.5, step=0.05, key="category_confidence")

            match_categories_btn = st.button("Match to Categories", type="primary", key="category_match")

        with col2:
            st.subheader("Results")
            results_placeholder_category = st.empty()
            if match_categories_btn:
                if st.session_state.category_input:
                    results_html = categorize_products_by_category(
                        st.session_state.category_input,
                        False,
                        category_use_expansion,
                        category_top_n,
                        category_confidence
                    )
                    results_placeholder_category.markdown(results_html, unsafe_allow_html=True)
                else:
                    results_placeholder_category.warning("Please enter product names.")

    # --- Common function for Reranking Tabs ---
    def create_reranking_ui(tab, tab_key_prefix, tab_name, backend_function, default_match="categories"):
        with tab:
            st.header(f"Match using {tab_name}")
            col1, col2 = st.columns(2)
            with col1:
                if st.button("Load Examples", key=f"{tab_key_prefix}_examples"):
                    st.session_state[f"{tab_key_prefix}_input"] = load_examples()

                tab_input_value = st.text_area(
                    "Product Names (one per line)",
                    value=st.session_state[f"{tab_key_prefix}_input"],
                    placeholder="Enter product names, one per line",
                    height=250,
                    key=f"{tab_key_prefix}_input_widget"
                )
                st.session_state[f"{tab_key_prefix}_input"] = tab_input_value # Update state

                tab_expansion = st.checkbox(
                    "Use Description Expansion (AI)",
                    value=False,
                    key=f"{tab_key_prefix}_expansion",
                    help="Expand product descriptions using AI before matching"
                )
                tab_emb_top_n = st.slider("Embedding Top N Results", 1, 50, 20, step=1, key=f"{tab_key_prefix}_emb_top_n")
                tab_top_n = st.slider("Final Top N Results", 1, 10, 5, step=1, key=f"{tab_key_prefix}_final_top_n")
                tab_confidence = st.slider("Matching Threshold", 0.1, 0.9, 0.5, step=0.05, key=f"{tab_key_prefix}_confidence")
                tab_match_type = st.radio(
                    "Match Type",
                    options=["categories", "ingredients"],
                    index=0 if default_match == "categories" else 1,
                    key=f"{tab_key_prefix}_match_type",
                    horizontal=True,
                    help="Choose whether to match against ingredients or categories"
                )

                tab_match_btn = st.button(f"Match using {tab_name}", type="primary", key=f"{tab_key_prefix}_match")

            with col2:
                st.subheader("Results")
                results_placeholder_rerank = st.empty()
                if tab_match_btn:
                    if st.session_state[f"{tab_key_prefix}_input"]:
                        results_html = backend_function(
                            st.session_state[f"{tab_key_prefix}_input"],
                            False,
                            tab_expansion,
                            tab_emb_top_n,
                            tab_top_n,
                            tab_confidence,
                            tab_match_type
                        )
                        results_placeholder_rerank.markdown(results_html, unsafe_allow_html=True)
                    else:
                        results_placeholder_rerank.warning("Please enter product names.")

    # Create the reranking tabs
    create_reranking_ui(tab_voyage, "voyage", "Voyage AI Reranking", categorize_products_with_voyage_reranking, "categories")
    create_reranking_ui(tab_openai, "openai", "OpenAI Reranking", categorize_products_with_openai_reranking, "categories")

    # --- Compare Methods Tab ---
    with tab_compare:
        st.header("Compare Matching Methods")
        col1, col2 = st.columns(2)
        with col1:
            if st.button("Load Examples", key="compare_examples"):
                st.session_state.compare_input = load_examples()

            compare_product_input_value = st.text_area(
                "Product Names (one per line)",
                value=st.session_state.compare_input,
                placeholder="4 Tbsp sweet pickle relish\nchocolate chips\nfresh parsley",
                height=200,
                key="compare_input_widget"
            )
            st.session_state.compare_input = compare_product_input_value # Update state

            compare_embedding_top_n = st.slider(
                "Initial embedding candidates",
                min_value=5, max_value=50, value=20, step=5,
                key="compare_emb_top_n"
            )
            compare_final_top_n = st.slider(
                "Final results per method",
                min_value=1, max_value=10, value=3, step=1,
                key="compare_final_top_n"
            )
            compare_confidence_threshold = st.slider(
                "Confidence threshold",
                min_value=0.0, max_value=1.0, value=0.5, step=0.05,
                key="compare_confidence"
            )
            compare_match_type = st.radio(
                "Match Type",
                options=["categories", "ingredients"],
                index=0,
                key="compare_match_type",
                horizontal=True,
                help="Choose whether to match against ingredients or categories"
            )
            compare_expansion = st.checkbox(
                "Use Description Expansion (AI)",
                value=False,
                key="compare_expansion",
                help="Expand product descriptions using AI before matching"
            )

            compare_btn = st.button("Compare Methods", type="primary", key="compare_run")

        with col2:
            st.subheader("Comparison Results")
            results_placeholder_compare = st.empty()
            if compare_btn:
                if st.session_state.compare_input:
                    results_html = compare_ingredient_methods_ui(
                        st.session_state.compare_input,
                        compare_embedding_top_n,
                        compare_final_top_n,
                        compare_confidence_threshold,
                        compare_match_type,
                        compare_expansion
                    )
                    results_placeholder_compare.markdown(results_html, unsafe_allow_html=True)
                else:
                    results_placeholder_compare.warning("Please enter product names.")

    st.markdown("---")
    st.markdown("Powered by Voyage AI embeddings • Built with Streamlit")