"""
Database explorer UI component for the pharmaceutical data management agent.
"""

import streamlit as st
import pandas as pd

def render_db_explorer_tab(session_state):
    """
    Render the database explorer tab in the UI.
    
    Args:
        session_state: Streamlit session state
    """
    st.subheader("Database Explorer")
    
    # Get tables by category
    tables = session_state.db.get_tables()
    
    # Display tables by category
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("### Raw Data Tables")
        for table in tables["raw_tables"]:
            with st.expander(table):
                sample = session_state.db.get_table_sample(table, 3)
                st.dataframe(pd.DataFrame(sample))
        
        st.markdown("### Staging Tables")
        for table in tables["staging_tables"]:
            with st.expander(table):
                sample = session_state.db.get_table_sample(table, 3)
                st.dataframe(pd.DataFrame(sample))
    
    with col2:
        st.markdown("### Analytics Ready Data")
        for table in tables["ard_tables"]:
            with st.expander(table):
                sample = session_state.db.get_table_sample(table, 3)
                st.dataframe(pd.DataFrame(sample))
        
        st.markdown("### Data Products")
        for table in tables["data_products"]:
            with st.expander(table):
                sample = session_state.db.get_table_sample(table, 3)
                st.dataframe(pd.DataFrame(sample))
    
    # SQL Query Executor
    st.markdown("### Query Explorer")
    with st.form(key="sql_form"):
        sql_query = st.text_area("Enter SQL Query", height=100, 
                                placeholder="SELECT * FROM ARD_SALES_PERFORMANCE WHERE region = 'North' LIMIT 5")
        run_sql = st.form_submit_button("Run Query")
    
    if run_sql and sql_query:
        with st.spinner("Executing query..."):
            result = session_state.db.execute_query(sql_query)
            
            if "error" in result:
                st.error(f"Error executing query: {result['error']}")
            elif "data" in result:
                st.dataframe(pd.DataFrame(result["data"]))
                st.success(f"Query returned {len(result['data'])} rows")
            elif "tables" in result:
                st.write(result["tables"])
            elif "schema" in result:
                st.write(f"Schema for {result['table']}:")
                st.dataframe(pd.DataFrame(result["schema"]))