import streamlit as st
from io import BytesIO
from datetime import datetime
from markdown import markdown
from bs4 import BeautifulSoup
from docx import Document

# Assuming these are correctly defined and accessible
from utils.pdf_utils import extract_text
from clients.azure_client import AzureClient
from clients.gemini_client import GeminiClient

from agents.tariff_agent import TariffTradeAgent
from agents.MonthlyPerformanceAgent import MonthlyPerformanceAgent
from agents.ConsumerHealthAgent import ConsumerHealthAgent
from agents.overall_agent import OverallAgent
from agents.sales_agent import SalesAgent
from agents.supply_chain_agent import SupplyChainAgent
from agents.enhanced_chat_agent import EnhancedChatAgent
from agents.qa_agent import QAAgent

# *** New agent imports ***
from agents.QTDPerformanceAgent import QTDPerformanceAgent
from agents.CategoryPerformanceAgent import CategoryPerformanceAgent
from agents.CompetitiveLandscapeAgent import CompetitiveLandscapeAgent
from agents.WeatherAgent import WeatherAgent
from agents.FY25ExpectationsAgent import FY25ExpectationsAgent
from agents.GeopoliticalUncertaintyAgent import GeopoliticalUncertaintyAgent
from agents.WholesaleAgent import WholesaleAgent
from agents.trf_agent import TariffAgent

def markdown_to_word(markdown_content, output_file='output.docx'):
    # Convert Markdown to HTML
    html_content = markdown(markdown_content)

    # Parse HTML using BeautifulSoup
    soup = BeautifulSoup(html_content, 'html.parser')

    # Create a Word document
    doc = Document()

    # Add content to the Word document
    for element in soup.descendants:
        if element.name == 'h1':
            doc.add_heading(element.text, level=1)
        elif element.name == 'h2':
            doc.add_heading(element.text, level=2)
        elif element.name == 'p':
            doc.add_paragraph(element.text)
        elif element.name == 'li':
            doc.add_paragraph(f"- {element.text}")

    # Save the Word document
    doc.save(output_file)


def initialize_clients_and_agents(provider_select: str, mode: str = None):
    need_reinit = (
        "client" not in st.session_state
        or st.session_state.get("provider") != provider_select
    )

    if need_reinit:
        if provider_select == "Azure GPT":
            st.session_state.client = AzureClient()
            is_azure = True
        else:
            st.session_state.client = GeminiClient()
            is_azure = False

        st.session_state.provider = provider_select

        # Dashboard agents
        st.session_state.agent_overall = OverallAgent(st.session_state.client, is_azure)
        st.session_state.agent_tariff = TariffTradeAgent(st.session_state.client, is_azure)
        st.session_state.agent_monthly_performance = MonthlyPerformanceAgent(st.session_state.client, is_azure)
        st.session_state.agent_consumer_health = ConsumerHealthAgent(st.session_state.client, is_azure)
        st.session_state.agent_sales = SalesAgent(st.session_state.client, is_azure)
        st.session_state.agent_supplychain = SupplyChainAgent(st.session_state.client, is_azure)

        # *** New agents ***
        st.session_state.agent_qtd_performance = QTDPerformanceAgent(st.session_state.client, is_azure)
        st.session_state.agent_category_performance = CategoryPerformanceAgent(st.session_state.client, is_azure)
        st.session_state.agent_competitive_landscape = CompetitiveLandscapeAgent(st.session_state.client, is_azure)
        st.session_state.agent_weather = WeatherAgent(st.session_state.client, is_azure)
        st.session_state.agent_fy25_expectations = FY25ExpectationsAgent(st.session_state.client, is_azure)
        st.session_state.agent_geopolitical_uncertainty = GeopoliticalUncertaintyAgent(st.session_state.client, is_azure)
        st.session_state.agent_wholesale = WholesaleAgent(st.session_state.client, is_azure)
        st.session_state.agent_trf = TariffAgent(st.session_state.client, is_azure)

        # Chat agent
        st.session_state.chat_agent = EnhancedChatAgent(st.session_state.client, is_azure)

        # QA agent if required
        if mode == "qa" or mode is None:
            st.session_state.agent_qa = QAAgent(st.session_state.client, is_azure)

    else:
        if mode in [None, "dashboard"]:
            is_azure = st.session_state.provider == "Azure GPT"
            if "agent_overall" not in st.session_state:
                st.session_state.agent_overall = OverallAgent(st.session_state.client, is_azure)
            if "agent_tariff" not in st.session_state:
                st.session_state.agent_tariff = TariffTradeAgent(st.session_state.client, is_azure)
            if "agent_trf" not in st.session_state:
                st.session_state.agent_trf = TariffAgent(st.session_state.client, is_azure)
            if "agent_sales" not in st.session_state:
                is_azure = st.session_state.provider == "Azure GPT"
                st.session_state.agent_sales = SalesAgent(st.session_state.client, is_azure)
            if "agent_supplychain" not in st.session_state:
                is_azure = st.session_state.provider == "Azure GPT"
                st.session_state.agent_supplychain = SupplyChainAgent(st.session_state.client, is_azure)
            if "agent_monthly_performance" not in st.session_state:
                st.session_state.agent_monthly_performance = MonthlyPerformanceAgent(st.session_state.client, is_azure)
            if "agent_consumer_health" not in st.session_state:
                st.session_state.agent_consumer_health = ConsumerHealthAgent(st.session_state.client, is_azure)

            # *** Initialize new agents if missing ***
            if "agent_qtd_performance" not in st.session_state:
                st.session_state.agent_qtd_performance = QTDPerformanceAgent(st.session_state.client, is_azure)
            if "agent_category_performance" not in st.session_state:
                st.session_state.agent_category_performance = CategoryPerformanceAgent(st.session_state.client, is_azure)
            if "agent_competitive_landscape" not in st.session_state:
                st.session_state.agent_competitive_landscape = CompetitiveLandscapeAgent(st.session_state.client, is_azure)
            if "agent_weather" not in st.session_state:
                st.session_state.agent_weather = WeatherAgent(st.session_state.client, is_azure)
            if "agent_fy25_expectations" not in st.session_state:
                st.session_state.agent_fy25_expectations = FY25ExpectationsAgent(st.session_state.client, is_azure)
            if "agent_geopolitical_uncertainty" not in st.session_state:
                st.session_state.agent_geopolitical_uncertainty = GeopoliticalUncertaintyAgent(st.session_state.client, is_azure)
            if "agent_wholesale" not in st.session_state:
                st.session_state.agent_wholesale = WholesaleAgent(st.session_state.client, is_azure)

            if "chat_agent" not in st.session_state:
                st.session_state.chat_agent = EnhancedChatAgent(st.session_state.client, is_azure)

        if mode in [None, "qa"]:
            if "agent_qa" not in st.session_state:
                is_azure = st.session_state.provider == "Azure GPT"
                st.session_state.agent_qa = QAAgent(st.session_state.client, is_azure)


def analyze_documents(mode, agent, processed_files, sess, custom_prompt=None):
    # Process individual files
    for filename, file_bytes in sess['uploaded_files'].items():
        if filename not in processed_files:
            try:
                file_obj = BytesIO(file_bytes)
                file_obj.name = filename
                text = extract_text(file_obj)
                if text:
                    # *** CHANGE: explicitly pass agent.prompt_analyze to ensure custom prompt used ***
                    topics = agent.analyze_document(
                        text,
                        custom_prompt=agent.prompt_analyze if hasattr(agent, 'prompt_analyze') else None
                    )

                    processed_files[filename] = {
                        "text": text,
                        "topics": topics,
                        "context": text,
                    }
                    st.toast(f"✅ Processed {filename}")
                else:
                    st.error(f"❌ Failed to extract text for {filename}")
            except Exception as e:
                st.error(f"Error processing {filename}: {str(e)}")

    # Generate consolidated summary
    summary = agent.generate_consolidated_summary_from_topics(processed_files, custom_prompt)

    if mode == "overall":
        sess['consolidated_summary_overall'] = summary
    elif mode == "tariff":
        sess['consolidated_summary_tariff'] = summary
    elif mode == "trf":
        sess['consolidated_summary_trf'] = summary
    elif mode == "sales":
        sess['consolidated_summary_sales'] = summary
    elif mode == "supplychain":
        sess['consolidated_summary_supplychain'] = summary
    elif mode == "monthly_performance":
        sess['consolidated_summary_monthly_performance'] = summary
    elif mode == "consumer_health":
        sess['consolidated_summary_consumer_health'] = summary
    elif mode == "qtd_performance":
        sess['consolidated_summary_qtd_performance'] = summary
    elif mode == "category_performance":
        sess['consolidated_summary_category_performance'] = summary
    elif mode == "competitive_landscape":
        sess['consolidated_summary_competitive_landscape'] = summary
    elif mode == "weather":
        sess['consolidated_summary_weather'] = summary
    elif mode == "fy25_expectations":
        sess['consolidated_summary_fy25_expectations'] = summary
    elif mode == "geopolitical_uncertainty":
        sess['consolidated_summary_geopolitical_uncertainty'] = summary
    elif mode == "wholesale":
        sess['consolidated_summary_wholesale'] = summary

    sess['last_analysis_mode'] = mode
    sess['active_file'] = None
    sess['active_topic'] = None
    st.success(f"Analysis and summary generated for {mode.replace('_', ' ').replace('consumer_health', 'consumer health').replace('monthly_performance', 'monthly performance')}.")


def main():
    st.set_page_config(page_title="Macy's Document Analysis Pro", layout="wide")

    st.sidebar.title("InsightHub")
    page = st.sidebar.radio("Select Any", ["Analysis Dashboard", "Question Assistant"])

    if page == "Analysis Dashboard":
        dashboard_page()
    elif page == "Question Assistant":
        qa_page()


def dashboard_page():
    st.title("📊 Macy's Document Analysis Dashboard")

    provider_select = st.radio("Select AI Provider", ["Azure GPT", "Google Gemini"])

    initialize_clients_and_agents(provider_select, mode="dashboard")

    if "session" not in st.session_state:
        st.session_state.session = {
            'uploaded_files': {},
            'processed_files_overall': {},
            'processed_files_tariff': {},
            'processed_files_trf': {},
            'processed_files_sales': {},
            'processed_files_supplychain': {},
            'processed_files_monthly_performance': {},
            'processed_files_consumer_health': {},
            'processed_files_qtd_performance': {},
            'processed_files_category_performance': {},
            'processed_files_competitive_landscape': {},
            'processed_files_weather': {},
            'processed_files_fy25_expectations': {},
            'processed_files_geopolitical_uncertainty': {},
            'processed_files_wholesale': {},
            'active_file': None,
            'active_topic': None,
            'topic_details_cache': {},
            'chat_history': [],
            'consolidated_summary_overall': None,
            'consolidated_summary_tariff': None,
            'consolidated_summary_trf': None,
            'consolidated_summary_sales': None,
            'consolidated_summary_supplychain': None,
            'consolidated_summary_monthly_performance': None,
            'consolidated_summary_consumer_health': None,
            'consolidated_summary_qtd_performance': None,
            'consolidated_summary_category_performance': None,
            'consolidated_summary_competitive_landscape': None,
            'consolidated_summary_weather': None,
            'consolidated_summary_fy25_expectations': None,
            'consolidated_summary_geopolitical_uncertainty': None,
            'consolidated_summary_wholesale': None,
            'last_analysis_mode': None,
            'custom_prompt_enabled': {},
            'custom_prompt_text': {}
        }

    sess = st.session_state.session

    # --- File Upload ---
    st.subheader("1. Upload Documents")
    uploaded_files = st.file_uploader(
        "Upload PDF documents",
        type=None,
        accept_multiple_files=True,
        key="file_uploader",
    )
    if uploaded_files:
        new_files_added = False
        for file in uploaded_files:
            if file.name not in sess['uploaded_files']:
                sess['uploaded_files'][file.name] = file.getvalue()
                new_files_added = True

        if new_files_added:
            with st.spinner("🔍 Analyzing new documents with Overall Agent..."):
                analyze_documents(
                    mode="overall",
                    agent=st.session_state.agent_overall,
                    processed_files=sess['processed_files_overall'],
                    sess=sess
                )

    # --- Custom Prompts & Analysis ---
    st.subheader("2. Custom Prompts & Analysis")

    if 'custom_prompts' not in sess:
        sess['custom_prompts'] = {
            'overall': {
                'summary': OverallAgent.prompt_summary,
                'analyze': OverallAgent.prompt_analyze,
                'detail': OverallAgent.prompt_detail
            },
            'tariff': {
                'summary': TariffTradeAgent.prompt_summary,
                'analyze': TariffTradeAgent.prompt_analyze,
                'detail': TariffTradeAgent.prompt_detail
            },
            'trf': {
                'summary': TariffAgent.prompt_summary,
                'analyze': TariffAgent.prompt_analyze,
                'detail': TariffAgent.prompt_detail
            },
            'sales': {
                'summary': SalesAgent.prompt_summary,
                'analyze': SalesAgent.prompt_analyze,
                'detail': SalesAgent.prompt_detail
            },
            'supplychain': {
                'summary': SupplyChainAgent.prompt_summary,
                'analyze': SupplyChainAgent.prompt_analyze,
                'detail': SupplyChainAgent.prompt_detail
            },
            'monthly_performance': {
                'summary': MonthlyPerformanceAgent.prompt_summary,
                'analyze': MonthlyPerformanceAgent.prompt_analyze,
                'detail': MonthlyPerformanceAgent.prompt_detail
            },
            'consumer_health': {
                'summary': ConsumerHealthAgent.prompt_summary,
                'analyze': ConsumerHealthAgent.prompt_analyze,
                'detail': ConsumerHealthAgent.prompt_detail
            },
            'qtd_performance': {
                'summary': QTDPerformanceAgent.prompt_summary,
                'analyze': QTDPerformanceAgent.prompt_analyze,
                'detail': QTDPerformanceAgent.prompt_detail
            },
            'category_performance': {
                'summary': CategoryPerformanceAgent.prompt_summary,
                'analyze': CategoryPerformanceAgent.prompt_analyze,
                'detail': CategoryPerformanceAgent.prompt_detail
            },
            'competitive_landscape': {
                'summary': CompetitiveLandscapeAgent.prompt_summary,
                'analyze': CompetitiveLandscapeAgent.prompt_analyze,
                'detail': CompetitiveLandscapeAgent.prompt_detail
            },
            'weather': {
                'summary': WeatherAgent.prompt_summary,
                'analyze': WeatherAgent.prompt_analyze,
                'detail': WeatherAgent.prompt_detail
            },
            'fy25_expectations': {
                'summary': FY25ExpectationsAgent.prompt_summary,
                'analyze': FY25ExpectationsAgent.prompt_analyze,
                'detail': FY25ExpectationsAgent.prompt_detail
            },
            'geopolitical_uncertainty': {
                'summary': GeopoliticalUncertaintyAgent.prompt_summary,
                'analyze': GeopoliticalUncertaintyAgent.prompt_analyze,
                'detail': GeopoliticalUncertaintyAgent.prompt_detail
            },
            'wholesale': {
                'summary': WholesaleAgent.prompt_summary,
                'analyze': WholesaleAgent.prompt_analyze,
                'detail': WholesaleAgent.prompt_detail
            }
        }

    tab_overall, tab_tariff, tab_trf,  tab_sales, tab_supply , tab_monthly_performance, tab_consumer_health, tab_qtd_performance, tab_category_performance, tab_competitive_landscape, tab_weather, tab_fy25_expectations, tab_geopolitical_uncertainty, tab_wholesale = st.tabs([
        "📑 Overall", "🔍 US Politics(Tariff)", "🔍 Tariff", "📈 Sales", "🚚 Supply Chain" , "📈 Monthly Performance", "🩺 Consumer Health",
        "📊 QTD Performance", "🛒 Category Performance", "🏆 Competitive Landscape", "☁️ Weather",
        "📅 FY25 Expectations", "🌍 Geopolitical Uncertainty", "🏬 Wholesale"
    ])

    generate_mode = None

    with tab_overall:
        st.subheader("Overall Prompts")
        use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_overall")

        if use_custom:
            sess['custom_prompts']['overall']['summary'] = st.text_area(
                "Summary Prompt:",
                value=sess['custom_prompts']['overall']['summary'],
                height=150,
                key="overall_prompt_summary"
            )
            sess['custom_prompts']['overall']['analyze'] = st.text_area(
                "Analysis Prompt:",
                value=sess['custom_prompts']['overall']['analyze'],
                height=150,
                key="overall_prompt_analyze"
            )
            sess['custom_prompts']['overall']['detail'] = st.text_area(
                "Detail Prompt:",
                value=sess['custom_prompts']['overall']['detail'],
                height=150,
                key="overall_prompt_detail"
            )

            if st.button("🔄 Generate with Custom Prompts", key="generate_overall"):
                generate_mode = "overall"

    with tab_tariff:
        st.subheader("US Politics Tariff Prompts")
        use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_tariff")

        if use_custom:
            sess['custom_prompts']['tariff']['summary'] = st.text_area(
                "Summary Prompt:",
                value=sess['custom_prompts']['tariff']['summary'],
                height=150,
                key="tariff_prompt_summary"
            )
            sess['custom_prompts']['tariff']['analyze'] = st.text_area(
                "Analysis Prompt:",
                value=sess['custom_prompts']['tariff']['analyze'],
                height=150,
                key="tariff_prompt_analyze"
            )
            sess['custom_prompts']['tariff']['detail'] = st.text_area(
                "Detail Prompt:",
                value=sess['custom_prompts']['tariff']['detail'],
                height=150,
                key="tariff_prompt_detail"
            )

            if st.button("🔄 Generate with Custom Prompts", key="generate_tariff"):
                generate_mode = "tariff"
    
    with tab_trf:
        st.subheader("Tariff Prompts")
        use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_trf")

        if use_custom:
            sess['custom_prompts']['trf']['summary'] = st.text_area(
                "Summary Prompt:",
                value=sess['custom_prompts']['trf']['summary'],
                height=150,
                key="trf_prompt_summary"
            )
            sess['custom_prompts']['trf']['analyze'] = st.text_area(
                "Analysis Prompt:",
                value=sess['custom_prompts']['trf']['analyze'],
                height=150,
                key="trf_prompt_analyze"
            )
            sess['custom_prompts']['trf']['detail'] = st.text_area(
                "Detail Prompt:",
                value=sess['custom_prompts']['trf']['detail'],
                height=150,
                key="trf_prompt_detail"
            )

            if st.button("🔄 Generate with Custom Prompts", key="generate_trf"):
                generate_mode = "trf"
    
    with tab_sales:
        st.subheader("Sales Prompts")
        use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_sales")

        if use_custom:
            sess['custom_prompts']['sales']['summary'] = st.text_area(
                "Summary Prompt:",
                value=sess['custom_prompts']['sales']['summary'],
                height=150,
                key="sales_prompt_summary"
            )
            sess['custom_prompts']['sales']['analyze'] = st.text_area(
                "Analysis Prompt:",
                value=sess['custom_prompts']['sales']['analyze'],
                height=150,
                key="sales_prompt_analyze"
            )
            sess['custom_prompts']['sales']['detail'] = st.text_area(
                "Detail Prompt:",
                value=sess['custom_prompts']['sales']['detail'],
                height=150,
                key="sales_prompt_detail"
            )

            if st.button("🔄 Generate with Custom Prompts", key="generate_sales"):
                generate_mode = "sales"

    with tab_supply:
        st.subheader("Supply Chain Prompts")
        use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_supply")

        if use_custom:
            sess['custom_prompts']['supplychain']['summary'] = st.text_area(
                "Summary Prompt:",
                value=sess['custom_prompts']['supplychain']['summary'],
                height=150,
                key="supply_prompt_summary"
            )
            sess['custom_prompts']['supplychain']['analyze'] = st.text_area(
                "Analysis Prompt:",
                value=sess['custom_prompts']['supplychain']['analyze'],
                height=150,
                key="supply_prompt_analyze"
            )
            sess['custom_prompts']['supplychain']['detail'] = st.text_area(
                "Detail Prompt:",
                value=sess['custom_prompts']['supplychain']['detail'],
                height=150,
                key="supply_prompt_detail"
            )

            if st.button("🔄 Generate with Custom Prompts", key="generate_supplychain"):
                generate_mode = "supplychain"

    with tab_monthly_performance:
        st.subheader("Monthly Performance Prompts")
        use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_monthly_performance")

        if use_custom:
            sess['custom_prompts']['monthly_performance']['summary'] = st.text_area(
                "Summary Prompt:",
                value=sess['custom_prompts']['monthly_performance']['summary'],
                height=150,
                key="monthly_performance_prompt_summary"
            )
            sess['custom_prompts']['monthly_performance']['analyze'] = st.text_area(
                "Analysis Prompt:",
                value=sess['custom_prompts']['monthly_performance']['analyze'],
                height=150,
                key="monthly_performance_prompt_analyze"
            )
            sess['custom_prompts']['monthly_performance']['detail'] = st.text_area(
                "Detail Prompt:",
                value=sess['custom_prompts']['monthly_performance']['detail'],
                height=150,
                key="monthly_performance_prompt_detail"
            )

            if st.button("🔄 Generate with Custom Prompts", key="generate_monthly_performance"):
                generate_mode = "monthly_performance"

    with tab_consumer_health:
        st.subheader("Consumer Health Prompts")
        use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_consumer_health")

        if use_custom:
            sess['custom_prompts']['consumer_health']['summary'] = st.text_area(
                "Summary Prompt:",
                value=sess['custom_prompts']['consumer_health']['summary'],
                height=150,
                key="consumer_health_prompt_summary"
            )
            sess['custom_prompts']['consumer_health']['analyze'] = st.text_area(
                "Analysis Prompt:",
                value=sess['custom_prompts']['consumer_health']['analyze'],
                height=150,
                key="consumer_health_prompt_analyze"
            )
            sess['custom_prompts']['consumer_health']['detail'] = st.text_area(
                "Detail Prompt:",
                value=sess['custom_prompts']['consumer_health']['detail'],
                height=150,
                key="consumer_health_prompt_detail"
            )

            if st.button("🔄 Generate with Custom Prompts", key="generate_consumer_health"):
                generate_mode = "consumer_health"

    with tab_qtd_performance:
        st.subheader("QTD Performance Prompts")
        use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_qtd_performance")

        if use_custom:
            sess['custom_prompts']['qtd_performance']['summary'] = st.text_area(
                "Summary Prompt:",
                value=sess['custom_prompts']['qtd_performance']['summary'],
                height=150,
                key="qtd_performance_prompt_summary"
            )
            sess['custom_prompts']['qtd_performance']['analyze'] = st.text_area(
                "Analysis Prompt:",
                value=sess['custom_prompts']['qtd_performance']['analyze'],
                height=150,
                key="qtd_performance_prompt_analyze"
            )
            sess['custom_prompts']['qtd_performance']['detail'] = st.text_area(
                "Detail Prompt:",
                value=sess['custom_prompts']['qtd_performance']['detail'],
                height=150,
                key="qtd_performance_prompt_detail"
            )

            if st.button("🔄 Generate with Custom Prompts", key="generate_qtd_performance"):
                generate_mode = "qtd_performance"

    with tab_category_performance:
        st.subheader("Category Performance Prompts")
        use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_category_performance")

        if use_custom:
            sess['custom_prompts']['category_performance']['summary'] = st.text_area(
                "Summary Prompt:",
                value=sess['custom_prompts']['category_performance']['summary'],
                height=150,
                key="category_performance_prompt_summary"
            )
            sess['custom_prompts']['category_performance']['analyze'] = st.text_area(
                "Analysis Prompt:",
                value=sess['custom_prompts']['category_performance']['analyze'],
                height=150,
                key="category_performance_prompt_analyze"
            )
            sess['custom_prompts']['category_performance']['detail'] = st.text_area(
                "Detail Prompt:",
                value=sess['custom_prompts']['category_performance']['detail'],
                height=150,
                key="category_performance_prompt_detail"
            )

            if st.button("🔄 Generate with Custom Prompts", key="generate_category_performance"):
                generate_mode = "category_performance"

    with tab_competitive_landscape:
        st.subheader("Competitive Landscape Prompts")
        use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_competitive_landscape")

        if use_custom:
            sess['custom_prompts']['competitive_landscape']['summary'] = st.text_area(
                "Summary Prompt:",
                value=sess['custom_prompts']['competitive_landscape']['summary'],
                height=150,
                key="competitive_landscape_prompt_summary"
            )
            sess['custom_prompts']['competitive_landscape']['analyze'] = st.text_area(
                "Analysis Prompt:",
                value=sess['custom_prompts']['competitive_landscape']['analyze'],
                height=150,
                key="competitive_landscape_prompt_analyze"
            )
            sess['custom_prompts']['competitive_landscape']['detail'] = st.text_area(
                "Detail Prompt:",
                value=sess['custom_prompts']['competitive_landscape']['detail'],
                height=150,
                key="competitive_landscape_prompt_detail"
            )

            if st.button("🔄 Generate with Custom Prompts", key="generate_competitive_landscape"):
                generate_mode = "competitive_landscape"

    with tab_weather:
        st.subheader("Weather Prompts")
        use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_weather")

        if use_custom:
            sess['custom_prompts']['weather']['summary'] = st.text_area(
                "Summary Prompt:",
                value=sess['custom_prompts']['weather']['summary'],
                height=150,
                key="weather_prompt_summary"
            )
            sess['custom_prompts']['weather']['analyze'] = st.text_area(
                "Analysis Prompt:",
                value=sess['custom_prompts']['weather']['analyze'],
                height=150,
                key="weather_prompt_analyze"
            )
            sess['custom_prompts']['weather']['detail'] = st.text_area(
                "Detail Prompt:",
                value=sess['custom_prompts']['weather']['detail'],
                height=150,
                key="weather_prompt_detail"
            )

            if st.button("🔄 Generate with Custom Prompts", key="generate_weather"):
                generate_mode = "weather"

    with tab_fy25_expectations:
        st.subheader("FY25 Expectations Prompts")
        use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_fy25_expectations")

        if use_custom:
            sess['custom_prompts']['fy25_expectations']['summary'] = st.text_area(
                "Summary Prompt:",
                value=sess['custom_prompts']['fy25_expectations']['summary'],
                height=150,
                key="fy25_expectations_prompt_summary"
            )
            sess['custom_prompts']['fy25_expectations']['analyze'] = st.text_area(
                "Analysis Prompt:",
                value=sess['custom_prompts']['fy25_expectations']['analyze'],
                height=150,
                key="fy25_expectations_prompt_analyze"
            )
            sess['custom_prompts']['fy25_expectations']['detail'] = st.text_area(
                "Detail Prompt:",
                value=sess['custom_prompts']['fy25_expectations']['detail'],
                height=150,
                key="fy25_expectations_prompt_detail"
            )

            if st.button("🔄 Generate with Custom Prompts", key="generate_fy25_expectations"):
                generate_mode = "fy25_expectations"

    with tab_geopolitical_uncertainty:
        st.subheader("Geopolitical Uncertainty Prompts")
        use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_geopolitical_uncertainty")

        if use_custom:
            sess['custom_prompts']['geopolitical_uncertainty']['summary'] = st.text_area(
                "Summary Prompt:",
                value=sess['custom_prompts']['geopolitical_uncertainty']['summary'],
                height=150,
                key="geopolitical_uncertainty_prompt_summary"
            )
            sess['custom_prompts']['geopolitical_uncertainty']['analyze'] = st.text_area(
                "Analysis Prompt:",
                value=sess['custom_prompts']['geopolitical_uncertainty']['analyze'],
                height=150,
                key="geopolitical_uncertainty_prompt_analyze"
            )
            sess['custom_prompts']['geopolitical_uncertainty']['detail'] = st.text_area(
                "Detail Prompt:",
                value=sess['custom_prompts']['geopolitical_uncertainty']['detail'],
                height=150,
                key="geopolitical_uncertainty_prompt_detail"
            )

            if st.button("🔄 Generate with Custom Prompts", key="generate_geopolitical_uncertainty"):
                generate_mode = "geopolitical_uncertainty"

    with tab_wholesale:
        st.subheader("Wholesale Prompts")
        use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_wholesale")

        if use_custom:
            sess['custom_prompts']['wholesale']['summary'] = st.text_area(
                "Summary Prompt:",
                value=sess['custom_prompts']['wholesale']['summary'],
                height=150,
                key="wholesale_prompt_summary"
            )
            sess['custom_prompts']['wholesale']['analyze'] = st.text_area(
                "Analysis Prompt:",
                value=sess['custom_prompts']['wholesale']['analyze'],
                height=150,
                key="wholesale_prompt_analyze"
            )
            sess['custom_prompts']['wholesale']['detail'] = st.text_area(
                "Detail Prompt:",
                value=sess['custom_prompts']['wholesale']['detail'],
                height=150,
                key="wholesale_prompt_detail"
            )

            if st.button("🔄 Generate with Custom Prompts", key="generate_wholesale"):
                generate_mode = "wholesale"

    st.subheader("Quick Analysis")

    # Buttons in 3 rows of 4,4 and 3
    cols1 = st.columns(4)
    analyze_overall = cols1[0].button("📑 Overall Summary")
    analyze_tariff = cols1[1].button("🔍 Analyze US Politics / Tariffs")
    analyze_monthly_performance = cols1[2].button("📈 Analyze Monthly Performance")
    analyze_consumer_health = cols1[3].button("🩺 Analyze Consumer Health")

    cols2 = st.columns(4)
    analyze_qtd_performance = cols2[0].button("📊 Analyze QTD Performance")
    analyze_category_performance = cols2[1].button("🛒 Analyze Category Performance")
    analyze_competitive_landscape = cols2[2].button("🏆 Analyze Competitive Landscape")
    analyze_weather = cols2[3].button("☁️ Analyze Weather")

    cols3 = st.columns(3)
    analyze_fy25_expectations = cols3[0].button("📅 Analyze FY25 Expectations")
    analyze_geopolitical_uncertainty = cols3[1].button("🌍 Analyze Geopolitical Uncertainty")
    analyze_wholesale = cols3[2].button("🏬 Analyze Wholesale")

    cols4 = st.columns(3)
    analyze_sales = cols4[0].button("📈 Analyze Sales")
    analyze_supplychain = cols4[1].button("🚚 Supply Chain")
    analyze_trf = cols4[2].button("🔍 Analyze Tariffs")
    

    if generate_mode or analyze_overall or analyze_tariff or analyze_sales or analyze_supplychain or analyze_monthly_performance or analyze_consumer_health or analyze_qtd_performance or analyze_category_performance or analyze_competitive_landscape or analyze_weather or analyze_fy25_expectations or analyze_geopolitical_uncertainty or analyze_wholesale:
        if not sess['uploaded_files']:
            st.error("Please upload documents first.")
            st.stop()

        mode = generate_mode or (
            "overall" if analyze_overall else
            "tariff" if analyze_tariff else
            "trf" if analyze_trf else
            "sales" if analyze_sales else
            "supplychain" if analyze_supplychain else 
            "monthly_performance" if analyze_monthly_performance else
            "consumer_health" if analyze_consumer_health else
            "qtd_performance" if analyze_qtd_performance else
            "category_performance" if analyze_category_performance else
            "competitive_landscape" if analyze_competitive_landscape else
            "weather" if analyze_weather else
            "fy25_expectations" if analyze_fy25_expectations else
            "geopolitical_uncertainty" if analyze_geopolitical_uncertainty else
            "wholesale" if analyze_wholesale else None
        )

        agent = None

        if mode == "overall":
            agent = st.session_state.agent_overall
        elif mode == "tariff":
            agent = st.session_state.agent_tariff
        elif mode == "trf":
            agent = st.session_state.agent_trf
        elif mode == "sales":
            agent = st.session_state.agent_sales
        elif mode == "supplychain":
            agent = st.session_state.agent_supplychain
        elif mode == "monthly_performance":
            agent = st.session_state.agent_monthly_performance
        elif mode == "consumer_health":
            agent = st.session_state.agent_consumer_health
        elif mode == "qtd_performance":
            agent = st.session_state.agent_qtd_performance
        elif mode == "category_performance":
            agent = st.session_state.agent_category_performance
        elif mode == "competitive_landscape":
            agent = st.session_state.agent_competitive_landscape
        elif mode == "weather":
            agent = st.session_state.agent_weather
        elif mode == "fy25_expectations":
            agent = st.session_state.agent_fy25_expectations
        elif mode == "geopolitical_uncertainty":
            agent = st.session_state.agent_geopolitical_uncertainty
        elif mode == "wholesale":
            agent = st.session_state.agent_wholesale

        sess[f'processed_files_{mode}'] = {}
        processed_files = sess[f'processed_files_{mode}']

        # Assign prompts on generate_mode or default
        if generate_mode:
            agent.prompt_summary = sess['custom_prompts'][mode]['summary']
            agent.prompt_analyze = sess['custom_prompts'][mode]['analyze']
            agent.prompt_detail = sess['custom_prompts'][mode]['detail']
            st.toast(f"Using custom prompts for {mode.replace('_', ' ')} analysis")
        else:
            agent_class = {
                "overall": OverallAgent,
                "tariff": TariffTradeAgent,
                "trf" : TariffAgent,
                "sales" : SalesAgent,
                "supplychain" : SupplyChainAgent,
                "monthly_performance": MonthlyPerformanceAgent,
                "consumer_health": ConsumerHealthAgent,
                "qtd_performance": QTDPerformanceAgent,
                "category_performance": CategoryPerformanceAgent,
                "competitive_landscape": CompetitiveLandscapeAgent,
                "weather": WeatherAgent,
                "fy25_expectations": FY25ExpectationsAgent,
                "geopolitical_uncertainty": GeopoliticalUncertaintyAgent,
                "wholesale": WholesaleAgent,
            }[mode]

            agent.prompt_summary = agent_class.prompt_summary
            agent.prompt_analyze = agent_class.prompt_analyze
            agent.prompt_detail = agent_class.prompt_detail

        with st.spinner(f"🔍 Analyzing {mode.replace('_', ' ')} documents..."):
            analyze_documents(
                mode,
                agent,
                processed_files,
                sess
            )

    last_mode = sess['last_analysis_mode']
    if last_mode is None:
        st.info("Select an analysis type to begin.")
        return

    mode_config = {
        'overall': {
            'files': sess['processed_files_overall'],
            'summary': sess['consolidated_summary_overall'],
            'sidebar_header': "Overall Topics",
            'main_header': "🌟 Executive Summary",
            'agent': st.session_state.agent_overall
        },
        'tariff': {
            'files': sess['processed_files_tariff'],
            'summary': sess['consolidated_summary_tariff'],
            'sidebar_header': "US Politics/Tariff Topics",
            'main_header': "🌟 US Politics/Tariff Analysis",
            'agent': st.session_state.agent_tariff
        },
        'trf': {
            'files': sess['processed_files_tariff'],
            'summary': sess['consolidated_summary_tariff'],
            'sidebar_header': "Tariff Topics",
            'main_header': "🌟 Tariff Analysis",
            'agent': st.session_state.agent_trf
        },
        'sales': {
            'files': sess['processed_files_sales'],
            'summary': sess['consolidated_summary_sales'],
            'sidebar_header': "Sales Topics",
            'main_header': "🌟 Sales Analysis",
            'agent': st.session_state.agent_sales
        },
        'supplychain': {
            'files': sess['processed_files_supplychain'],
            'summary': sess['consolidated_summary_supplychain'],
            'sidebar_header': "Supply Chain Topics",
            'main_header': "🌟 Supply Chain Analysis",
            'agent': st.session_state.agent_supplychain
        },
        'monthly_performance': {
            'files': sess['processed_files_monthly_performance'],
            'summary': sess['consolidated_summary_monthly_performance'],
            'sidebar_header': "Monthly Performance Topics",
            'main_header': "🌟 Monthly Performance Analysis",
            'agent': st.session_state.agent_monthly_performance
        },
        'consumer_health': {
            'files': sess['processed_files_consumer_health'],
            'summary': sess['consolidated_summary_consumer_health'],
            'sidebar_header': "Consumer Health Topics",
            'main_header': "🌟 Consumer Health Analysis",
            'agent': st.session_state.agent_consumer_health
        },
        'qtd_performance': {
            'files': sess['processed_files_qtd_performance'],
            'summary': sess['consolidated_summary_qtd_performance'],
            'sidebar_header': "QTD Performance Topics",
            'main_header': "🌟 QTD Performance Analysis",
            'agent': st.session_state.agent_qtd_performance
        },
        'category_performance': {
            'files': sess['processed_files_category_performance'],
            'summary': sess['consolidated_summary_category_performance'],
            'sidebar_header': "Category Performance Topics",
            'main_header': "🌟 Category Performance Analysis",
            'agent': st.session_state.agent_category_performance
        },
        'competitive_landscape': {
            'files': sess['processed_files_competitive_landscape'],
            'summary': sess['consolidated_summary_competitive_landscape'],
            'sidebar_header': "Competitive Landscape Topics",
            'main_header': "🌟 Competitive Landscape Analysis",
            'agent': st.session_state.agent_competitive_landscape
        },
        'weather': {
            'files': sess['processed_files_weather'],
            'summary': sess['consolidated_summary_weather'],
            'sidebar_header': "Weather Topics",
            'main_header': "🌟 Weather Analysis",
            'agent': st.session_state.agent_weather
        },
        'fy25_expectations': {
            'files': sess['processed_files_fy25_expectations'],
            'summary': sess['consolidated_summary_fy25_expectations'],
            'sidebar_header': "FY25 Expectations Topics",
            'main_header': "🌟 FY25 Expectations Analysis",
            'agent': st.session_state.agent_fy25_expectations
        },
        'geopolitical_uncertainty': {
            'files': sess['processed_files_geopolitical_uncertainty'],
            'summary': sess['consolidated_summary_geopolitical_uncertainty'],
            'sidebar_header': "Geopolitical Uncertainty Topics",
            'main_header': "🌟 Geopolitical Uncertainty Analysis",
            'agent': st.session_state.agent_geopolitical_uncertainty
        },
        'wholesale': {
            'files': sess['processed_files_wholesale'],
            'summary': sess['consolidated_summary_wholesale'],
            'sidebar_header': "Wholesale Topics",
            'main_header': "🌟 Wholesale Analysis",
            'agent': st.session_state.agent_wholesale
        }
    }

    current_mode = mode_config[last_mode]

    # --- Sidebar Topics ---
    st.sidebar.header(current_mode['sidebar_header'])

    if last_mode == "overall":
        topics_by_file = {}
        for filename, filedata in current_mode['files'].items():
            for topic in filedata["topics"]:
                topics_by_file.setdefault(topic['name'], []).append({
                    'file': filename,
                    'score': topic['score']
                })

        sorted_topics = sorted(topics_by_file.items(),
                               key=lambda x: max(x[1], key=lambda y: y['score'])['score'],
                               reverse=True)
        before_view_more, after_view_more = sorted_topics[:3], sorted_topics[3:]

        for topic_name, occurrences in before_view_more:
            total_score = sum(occ['score'] for occ in occurrences)
            file_count = len(occurrences)
            btn_label = f"{topic_name} (📈 {total_score:.1f} | 📑 {file_count})"

            if st.sidebar.button(
                btn_label,
                key=f"overall_topic_{topic_name}",
                use_container_width=True,
            ):
                sess['active_topic'] = topic_name
                sess['active_file'] = None

        if after_view_more and st.sidebar.button('View More', key='1234', use_container_width=True):
            for topic_name, occurrences in after_view_more:
                total_score = sum(occ['score'] for occ in occurrences)
                file_count = len(occurrences)
                btn_label = f"{topic_name} (📈 {total_score:.1f} | 📑 {file_count})"

                if st.sidebar.button(
                    btn_label,
                    key=f"overall_topic_more_{topic_name}",
                    use_container_width=True,
                ):
                    sess['active_topic'] = topic_name
                    sess['active_file'] = None

    else:  # other modes: tariff, monthly_performance, consumer_health, qtd, etc.
        for filename, filedata in current_mode['files'].items():
            expanded = (filename == sess['active_file'])
            with st.sidebar.expander(f"📄 {filename}", expanded=expanded):
                topics = sorted(filedata["topics"], key=lambda t: t["score"], reverse=True)
                for topic in topics[:3]:
                    is_selected = (
                        sess['active_file'] == filename and sess['active_topic'] == topic['name']
                    )
                    style = "primary" if is_selected else "secondary"

                    if st.button(
                        f"• {topic['name']} (Score: {topic['score']})",
                        key=f"{last_mode}_topic_{filename}_{topic['name']}",
                        type=style,
                        use_container_width=True,
                    ):
                        sess['active_file'] = filename
                        sess['active_topic'] = topic['name']

                if len(topics) > 3 and st.sidebar.button(f'View More for {filename}', key=f'{filename}_view_more', use_container_width=True):
                    for topic in topics[3:]:
                        is_selected = (
                            sess['active_file'] == filename and sess['active_topic'] == topic['name']
                        )
                        style = "primary" if is_selected else "secondary"

                        if st.button(
                            f"• {topic['name']} (Score: {topic['score']})",
                            key=f"{last_mode}_topic_more_{filename}_{topic['name']}",
                            type=style,
                            use_container_width=True,
                        ):
                            sess['active_file'] = filename
                            sess['active_topic'] = topic['name']

    # --- Topic Detail Analysis ---
    st.subheader("Detailed Analysis")
    if last_mode == "overall" and sess['active_topic']:
        mentions = []
        for filename, filedata in current_mode['files'].items():
            for topic in filedata["topics"]:
                if topic['name'] == sess['active_topic']:
                    mentions.append({
                        'file': filename,
                        'summary': topic.get('summary', ''),
                        'score': topic['score'],
                        'context': filedata['context']
                    })

        with st.expander("📊 Topic Prevalence"):
            cols = st.columns(3)
            cols[0].metric("Total Documents", len(mentions))
            avg_relevance = (sum(m['score'] for m in mentions) / len(mentions)) if mentions else 0
            cols[1].metric("Average Relevance", f"{avg_relevance:.1f}")
            cols[2].metric("First Mentioned", mentions[0]['file'] if mentions else "N/A")

        if mentions:
            combined_context = "\n\n".join(f"=== {m['file']} ===\n{m['context']}" for m in mentions)
            key_cache = (last_mode, sess['active_topic'])
            if key_cache in sess['topic_details_cache']:
                st.markdown(sess['topic_details_cache'][key_cache])
            else:
                placeholder = st.empty()
                collected_text = ""
                try:
                    for chunk in current_mode['agent'].generate_topic_detail(combined_context, sess['active_topic']):
                        collected_text = chunk
                        placeholder.markdown(collected_text + (" ▌" if "▌" in collected_text else ""))
                    sess['topic_details_cache'][key_cache] = collected_text
                    markdown_to_word(collected_text, sess['active_topic']+'.docx')
                    with open(sess['active_topic']+'.docx', 'rb') as f:
                        st.download_button("Download Doc", f, file_name=sess['active_topic']+'.docx')

                except Exception as e:
                    st.error(f"Analysis error: {str(e)}")

    elif sess['active_file'] and sess['active_topic']:
        key_cache = (last_mode, sess['active_file'], sess['active_topic'])
        if key_cache in sess['topic_details_cache']:
            st.markdown(sess['topic_details_cache'][key_cache])
        else:
            placeholder = st.empty()
            collected_text = ""
            try:
                for chunk in current_mode['agent'].generate_topic_detail(
                    current_mode['files'][sess['active_file']]['context'],
                    sess['active_topic'],
                ):
                    collected_text = chunk
                    placeholder.markdown(collected_text + (" ▌" if "▌" in collected_text else ""))
                sess['topic_details_cache'][key_cache] = collected_text
                markdown_to_word(collected_text, sess['active_topic']+'.docx')
                with open(sess['active_topic']+'.docx', 'rb') as f:
                    st.download_button("Download Doc", f, file_name=sess['active_topic']+'.docx')

            except Exception as e:
                st.error(f"Analysis error: {str(e)}")
    else:
        st.info("Select a topic from the sidebar")

    st.markdown("---")

    st.header(current_mode['main_header'])
    if current_mode['summary']:
        st.markdown(current_mode['summary'])
        markdown_to_word(current_mode['summary'], current_mode['sidebar_header']+'.docx')
        with open(current_mode['sidebar_header']+'.docx', 'rb') as f:
            st.download_button("Download Doc", f, file_name=current_mode['sidebar_header']+'.docx')

    else:
        st.info("Summary not available yet")

    st.divider()

    # --- Chat Interface ---
    st.subheader("💬 Document Chat")

    available_files = list(current_mode['files'].keys())
    selected_files = st.multiselect(
        "Select documents for context:",
        available_files,
        key=f"chat_selector_{last_mode}"
    )

    for msg in sess['chat_history']:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])

    question = st.chat_input("Ask about the documents...")
    if question:
        if not selected_files:
            st.error("Please select documents first")
            st.stop()

        sess['chat_history'].append({
            "role": "user",
            "content": question,
            "timestamp": datetime.utcnow().isoformat(),
        })

        try:
            context = "\n\n".join([
                f"=== {fname} ===\n{current_mode['files'][fname]['context']}"
                for fname in selected_files
            ])

            with st.chat_message("assistant"):
                response_placeholder = st.empty()
                full_response = ""
                q_type = None

                for chunk, q_type in st.session_state.chat_agent.generate_answer(question, context):
                    full_response = chunk
                    response_placeholder.markdown(full_response + " ▌")

                response_placeholder.markdown(full_response)

                sess['chat_history'].append({
                    "role": "assistant",
                    "content": full_response,
                    "timestamp": datetime.utcnow().isoformat(),
                    "analysis_type": q_type
                })

        except Exception as e:
            st.error(f"Chat error: {str(e)}")
            sess['chat_history'].append({
                "role": "assistant",
                "content": f"Error: {str(e)}",
                "timestamp": datetime.utcnow().isoformat(),
            })

    if st.button("🔄 Reset Session", type="secondary"):
        preserved_keys = [
            "client", "agent_overall", "agent_tariff", "agent_trf", "agent_sales", "agent_supplychain",
            "agent_monthly_performance", "agent_consumer_health", "chat_agent", "provider",
            "agent_qtd_performance", "agent_category_performance", "agent_competitive_landscape",
            "agent_weather", "agent_fy25_expectations", "agent_geopolitical_uncertainty", "agent_wholesale"
        ]
        preserved = {k: st.session_state[k] for k in preserved_keys if k in st.session_state}

        st.session_state.clear()
        st.session_state.update(preserved)

        st.session_state.session = {
            'uploaded_files': {},
            'processed_files_overall': {},
            'processed_files_tariff': {},
            'processed_files_trf': {},
            'processed_files_monthly_performance': {},
            'processed_files_consumer_health': {},
            'processed_files_qtd_performance': {},
            'processed_files_category_performance': {},
            'processed_files_competitive_landscape': {},
            'processed_files_weather': {},
            'processed_files_fy25_expectations': {},
            'processed_files_geopolitical_uncertainty': {},
            'processed_files_wholesale': {},
            'active_file': None,
            'active_topic': None,
            'topic_details_cache': {},
            'chat_history': [],
            'consolidated_summary_overall': None,
            'consolidated_summary_tariff': None,
            'consolidated_summary_trf': None,
            'consolidated_summary_sales': None,
            'consolidated_summary_supplychain': None,
            'consolidated_summary_monthly_performance': None,
            'consolidated_summary_consumer_health': None,
            'consolidated_summary_qtd_performance': None,
            'consolidated_summary_category_performance': None,
            'consolidated_summary_competitive_landscape': None,
            'consolidated_summary_weather': None,
            'consolidated_summary_fy25_expectations': None,
            'consolidated_summary_geopolitical_uncertainty': None,
            'consolidated_summary_wholesale': None,
            'last_analysis_mode': None,
            'custom_prompt_enabled': {},
            'custom_prompt_text': {}
        }

        st.experimental_rerun()


def qa_page():
    st.title("❓ Macy's Question Assistant")

    if 'qa_page' not in st.session_state:
        st.session_state.qa_page = {
            'uploaded_files': {},
            'processed_files': {},
            'active_file': None,
            'active_topic': None,
            'chat_history': [],
            'selected_files': []
        }
    qa_state = st.session_state.qa_page

    provider_select = st.sidebar.radio("Select AI Provider", ["Azure GPT", "Google Gemini"])

    initialize_clients_and_agents(provider_select, mode="qa")

    st.subheader("Upload Documents")
    uploaded_files = st.file_uploader(
        "Upload PDF documents for Q&A extraction",
        type=["pdf"],
        accept_multiple_files=True,
        key="qa_file_uploader"
    )

    if uploaded_files:
        new_files_added = False
        for file in uploaded_files:
            if file.name not in qa_state['uploaded_files']:
                qa_state['uploaded_files'][file.name] = file.getvalue()
                new_files_added = True

        if new_files_added:
            with st.spinner("🔍 Extracting Q&A from new documents..."):
                for filename in qa_state['uploaded_files']:
                    if filename not in qa_state['processed_files']:
                        try:
                            file_bytes = qa_state['uploaded_files'][filename]
                            file_obj = BytesIO(file_bytes)
                            file_obj.name = filename
                            text = extract_text(file_obj)
                            if text:
                                qa_data = st.session_state.agent_qa.analyze_question(text)
                                qa_state['processed_files'][filename] = {
                                    'text': text,
                                    'qa_data': qa_data
                                }
                                st.toast(f"✅ Processed {filename}")

                                if not qa_state.get('error') and qa_data.get('topics'):
                                    qa_state['active_topic'] = qa_data['topics'][0]
                            else:
                                st.error(f"❌ Failed to extract text for {filename}")
                        except Exception as e:
                            st.error(f"Error processing {filename}: {str(e)}")

            if not qa_state['active_file'] and qa_state['processed_files']:
                qa_state['active_file'] = next(iter(qa_state['processed_files'].keys()))

    st.sidebar.header("📄 Processed Documents")
    for filename in qa_state['processed_files']:
        if st.sidebar.button(
            f"📄 {filename}",
            key=f"qa_file_{filename}",
            use_container_width=True,
        ):
            qa_state['active_file'] = filename
            file_data = qa_state['processed_files'][filename]
            qa_data = file_data['qa_data']
            if not qa_data.get('error') and qa_data.get('topics'):
                qa_state['active_topic'] = qa_data['topics'][0]

    if qa_state['active_file']:
        file_data = qa_state['processed_files'][qa_state['active_file']]
        qa_data = file_data['qa_data']

        st.sidebar.header("📌 Topics")

        if "error" in qa_data:
            st.sidebar.error("Error in document processing")

        elif 'topics' in qa_data:
            topics = qa_data['topics']

            for topic in topics:
                btn_label = f"{topic['topic']} ({topic['mention_count']} mentions)"
                if st.sidebar.button(
                    btn_label,
                    key=f"topic_{topic['topic']}_{qa_state['active_file']}",
                    use_container_width=True,
                ):
                    qa_state['active_topic'] = topic
        else:
            st.sidebar.info("No topics found in this document")

    st.subheader("📝 Extracted Q&A Content")

    if qa_state['active_file']:
        file_data = qa_state['processed_files'][qa_state['active_file']]
        qa_data = file_data['qa_data']
        st.markdown(f"**Document:** {qa_state['active_file']}")

        if "error" in qa_data:
            st.error(f"Error processing document: {qa_data['error']}")
            if "raw_response" in qa_data:
                st.text_area("Raw AI Response", qa_data["raw_response"], height=300)

        elif 'key_takeaways' in qa_data and qa_data['key_takeaways']:
            st.subheader("Key Takeaways")
            for takeaway in qa_data['key_takeaways']:
                st.markdown(f"- {takeaway}")

        if qa_state['active_topic']:
            topic = qa_state['active_topic']
            st.divider()
            st.subheader(topic['topic'])
            st.markdown(f"**Mentioned:** {topic['mention_count']} times")
            st.divider()

            if 'questions' in topic and topic['questions']:
                for idx, question in enumerate(topic['questions']):
                    firm = question.get('firm', 'Unknown Firm')
                    st.markdown(f"**Firm:** {firm}")
                    st.markdown(f"**Question {idx+1}:** {question['question']}")
                    st.write("---")
            else:
                st.info("No questions found for this topic")
        elif 'topics' in qa_data and qa_data['topics']:
            qa_state['active_topic'] = qa_data['topics'][0]
            st.experimental_rerun()
        else:
            st.info("No Q&A content found for this document")
    else:
        st.info("Upload documents to see extracted Q&A content")

    st.divider()

    st.subheader("💬 Chat with Documents")

    available_files = list(qa_state['processed_files'].keys())
    selected_files = st.multiselect(
        "Select documents to include in chat context:",
        available_files,
        default=available_files[:min(3, len(available_files))],
        key="qa_chat_selector"
    )
    qa_state['selected_files'] = selected_files

    for msg in qa_state['chat_history']:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])

    question = st.chat_input("Ask about the documents...")
    if question:
        if not selected_files:
            st.error("Please select at least one document for context")
            st.stop()

        qa_state['chat_history'].append({
            "role": "user",
            "content": question,
            "timestamp": datetime.utcnow().isoformat(),
        })

        try:
            context = "\n\n".join([
                f"=== {filename} ===\n{qa_state['processed_files'][filename]['text']}"
                for filename in selected_files
            ])

            with st.chat_message("assistant"):
                response_placeholder = st.empty()
                full_response = ""

                for chunk, _ in st.session_state.chat_agent.generate_answer(question, context):
                    full_response = chunk
                    response_placeholder.markdown(full_response + " ▌")

                response_placeholder.markdown(full_response)

                qa_state['chat_history'].append({
                    "role": "assistant",
                    "content": full_response,
                    "timestamp": datetime.utcnow().isoformat(),
                })

        except Exception as e:
            st.error(f"Chat error: {str(e)}")
            qa_state['chat_history'].append({
                "role": "assistant",
                "content": f"Error: {str(e)}",
                "timestamp": datetime.utcnow().isoformat(),
            })

    if st.button("🔄 Reset Q&A Session", key="reset_qa_page"):
        qa_state.clear()
        qa_state.update({
            'uploaded_files': {},
            'processed_files': {},
            'active_file': None,
            'active_topic': None,
            'chat_history': [],
            'selected_files': []
        })
        st.experimental_rerun()


if __name__ == "__main__":
    main()


when we click on tariff it is not picking the correct tarrif and getting me below errror it is picking us tariff insted

FileNotFoundError: [Errno 2] No such file or directory: 'US Politics/Tariff Topics.docx'


give me full correct code