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_trf 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_trf'],
            'summary': sess['consolidated_summary_trf'],
            '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()

this is my main.py

from utils.pdf_utils import parse_topics

class BaseAgent:
    def __init__(self, client, is_azure):
        self.client = client
        self.is_azure = is_azure

    def analyze_document(self, prompt_system, prompt_user, text):
        if self.is_azure:
            messages = [
                {"role": "system", "content": prompt_system},
                {"role": "user", "content": prompt_user + "\n\n" + text},
            ]
            response = self.client.generate(messages=messages, temperature=0)
            raw_text = response.choices[0].message.content.strip()
        else:
            raw_text = self.client.generate(prompt_system, prompt_user + "\n\n" + text)
        return parse_topics(raw_text)
    
    def analyze_question(self, prompt_system, prompt_user, text):
        """New method for QA extraction"""
        if self.is_azure:
            messages = [
                {"role": "system", "content": prompt_system},
                {"role": "user", "content": prompt_user + "\n\n" + text},
            ]
            response = self.client.generate(messages=messages, temperature=0)
            return response.choices[0].message.content.strip()
        else:
            return self.client.generate(prompt_system, prompt_user + "\n\n" + text)    

    def generate_consolidated_summary_from_topics(self, prompt_system, prompt_user_template, processed_files):
        topics_block = ""
        for filename, filedata in processed_files.items():
            topics_block += f"### Topics from document: {filename}\n"
            for topic in sorted(filedata["topics"], key=lambda x: x["score"], reverse=True):
                topics_block += f"- **{topic['name']}**: {topic['summary']} (Relevance: {topic['score']})\n"
            topics_block += "\n"

        prompt_user = prompt_user_template.format(topics_block=topics_block)

        if self.is_azure:
            messages = [
                {"role": "system", "content": prompt_system},
                {"role": "user", "content": prompt_user},
            ]
            response = self.client.generate(messages=messages, temperature=0)
            return response.choices[0].message.content.strip()
        else:
            return self.client.generate(prompt_system, prompt_user)

    def generate_topic_detail(self, prompt_system, prompt_user_template, document_text, topic):
        prompt_user = prompt_user_template.format(text=document_text, topic=topic)
        if self.is_azure:
            messages = [
                {"role": "system", "content": prompt_system},
                {"role": "user", "content": prompt_user},
            ]
            response = self.client.generate(messages=messages, temperature=0, stream=True)
            
            collected_chunks = []
            for chunk in response:
                if hasattr(chunk, "choices") and chunk.choices:
                    delta = chunk.choices[0].delta
                    content = getattr(delta, "content", "")
                    if content:
                        collected_chunks.append(content)
                        yield "".join(collected_chunks) + "▌"
            
            # Final yield without cursor
            if collected_chunks:
                yield "".join(collected_chunks)
        else:
            # For Gemini, we'll simulate streaming with 3 chunks
            full_response = self.client.generate(prompt_system, prompt_user)
            chunks = self.preserve_formatting_chunks(full_response)
            
            accumulated = ""
            for chunk in chunks:
                accumulated += chunk
                yield accumulated + "▌"
            
            yield accumulated  # Final yield without cursor
    
    def preserve_formatting_chunks(self, text):
        """Split response while preserving Markdown formatting"""
        # Split by paragraphs first
        paragraphs = text.split('\n\n')
        
        chunks = []
        for para in paragraphs:
            if para.startswith('##') or para.startswith('###') or para.startswith('* '):
                # Preserve headers and lists as whole chunks
                chunks.append(para + '\n\n')
            else:
                # Split long paragraphs into sentences
                sentences = para.split('. ')
                current_chunk = ""
                for sentence in sentences:
                    if not sentence:
                        continue
                    if len(current_chunk) + len(sentence) > 150:
                        if current_chunk:
                            chunks.append(current_chunk.strip() + '. ')
                        current_chunk = sentence
                    else:
                        current_chunk += sentence + '. '
                if current_chunk:
                    chunks.append(current_chunk.strip() + '\n\n')
        
        return chunks


            

this is base agent

from agents.base_agent import BaseAgent

class TariffTradeAgent(BaseAgent):
    prompt_sys = "Tariff and trade policy analyst for Macy's CEO works for only provided documents"
    
    # Universal formatting enforcement line
    formatting_directive = """Please strictly follow these formatting guidelines for all numerical mentions and sentence structures:
1. **No Combined Elements**: Always space-separate numbers, words, and units (e.g., "3.3 billion" not "3.3billion").
2. **Time Formatting**: Use formats like "Q4 2024" (never "Q42024" or "4th quarter").
3. **Number Rules**:  
   - 3.3B → 3.3 billion  
   - 150.2M → 150.2 million  
   - 0.33 → keep exact decimal  
4. **Text Integrity**:  
   - "net income" never "netincome"  
   - "year-over-year" never "YOY" or "yearoveryear"  
5. **Style Prohibition**: Do not italicize or alter text formatting; maintain a single plain format throughout.
6. **Bad Example**: `3.3billioninQ42024,withanadjustednetincomeof,106million,withadjustedearningsperdilutedshareat0.95` → REJECTED
7. **ALWAYS Required Format**: `3.3 billion in Q4 2024, with an adjusted net income of` → ACCEPTED

APPLY THESE RULES TO ALL NUMERICAL MENTIONS AND SENTENCE STRUCTURES."""
    
    prompt_analyze = (
        "You are a senior tariff and trade analyst for Macy's CEO.\n\n"
        "Carefully identify all significant tariff and trade policy topics impacting Macy's business operations, especially US Politics / Tariffs.\n\n"
        "For each topic, provide:\n"
        "- A concise topic name (e.g., \"US Politics / Tariffs\").\n"
        "- A brief summary in 100 words or less in third person (the company, management), covering: country or regional sourcing exposure where applicable, any quantitative financial impact if provided, and mitigation actions (such as pricing changes). If no details are available, say “N/A.”\n"
        "- A relevance score from 1 (minor) to 10 (critical).\n\n"
        "Format your response as one topic per line:\n"
        "- Topic Name: Summary [Relevance Score]\n\n"
        "Example:\n"
        "- US Politics / Tariffs: The company sources 40% of products from China, tariffs added 150 million in costs in Q3 2023, mitigation includes 3% price increases. [9]\n\n"
        "* ALWAYS provide verbatim source excerpts under the heading \"Ground Truth\" as follows:\n"
        "  - For each document where the topic appears, extract the exact unmodified sentence(s) or paragraph(s) from the document.\n"
        "  - Format each excerpt as: **Ground Truth:** \"Exact text from the document...\" (Document Name)\n"
        "  - If the topic is present in multiple documents, include separate excerpts.\n"
        "  - Do NOT paraphrase or alter the text.\n\n"
        + formatting_directive)
    
    prompt_summary = (
        "You are a senior tariff and trade analyst preparing a consolidated tariff and trade policy report for Macy's executive team from all documents provided.\n\n"
        "Identify and describe all impactful tariff, trade, and related government policy topics across the documents provided that could affect Macy's business operations including supply chain, procurement, retail sales, and technology.\n\n"
        "For each topic:\n"
        "- Provide a concise summary by combining relevant information from the documents.\n"
        "- Include relevant financial impacts, examples, figures, or timelines.\n"
        "- Provide three to five actionable strategic recommendations based on the analysis.\n\n"
        "For every topic, ALWAYS provide each excerpt mentioned under the heading \"Ground Truth\" as follows:\n"
        "- For each document where the topic appears, extract the exact, unmodified sentence(s) or paragraph(s) from the original document.\n"
        "- Format each excerpt as: **Ground Truth:** \"Exact text from the document...\" (Document Name)\n"
        "- If a topic is in multiple documents, include separate excerpts for each document with proper line breaks.\n"
        "- Do NOT paraphrase or alter any text.\n"
        "- If no exact excerpts are found, do NOT fabricate or modify content.\n\n"
        "Format the report in clear markdown with sections:\n"
        "- Executive Summary\n"
        "- Key Themes\n"
        "- Strategic Recommendations\n"
        "- Risk Assessment\n\n"
        "STRICTLY enforce plain text formatting:\n"
        + formatting_directive
    )
    
    prompt_detail = (
        "You are preparing a detailed briefing for Macy's CEO on the tariff and trade policy topic '{topic}', based on the provided external company document.\n\n"
        "Provide:\n"
        "1. A clear overview of the topic.\n"
        "2. Key insights and relevant excerpts.\n"
        "3. Important figures, dates, or timelines if available.\n"
        "4. A detailed analysis on how this topic could impact Macy's retail operations, supply chain, procurement, and consumer behavior.\n"
        "5. Three to five actionable strategic recommendations for Macy's leadership.\n\n"
        "For every topic, ALWAYS provide verbatim source excerpts under the heading \"Ground Truth\" as follows:\n"
        "- For each document where the topic appears, extract the exact, unmodified sentence(s) or paragraph(s) from the original document.\n"
        "- Format each excerpt as: **Ground Truth:** \"Exact text from the document...\" (Document Name)\n"
        "- If a topic is present in multiple documents, include a separate excerpt for each.\n"
        "- Do NOT paraphrase, summarize, or alter the text.\n"
        "- If no exact excerpts exist, do NOT fabricate or modify content.\n\n"
        "Document:\n"
        "{text}\n\n"
        "Topic: {topic}\n\n"
        "Please provide a structured briefing with clear paragraphs and bold topics.\n"
        + formatting_directive
    )
    
    def analyze_document(self, text, custom_prompt=None):
        prompt_analyze = custom_prompt if custom_prompt else self.prompt_analyze
        return super().analyze_document(self.prompt_sys, prompt_analyze, text)

    def generate_consolidated_summary_from_topics(self, processed_files, custom_prompt=None):
        prompt_summary = custom_prompt if custom_prompt else self.prompt_summary
        return super().generate_consolidated_summary_from_topics(self.prompt_sys, prompt_summary, processed_files)

    def generate_topic_detail(self, text, topic, custom_prompt=None):
        prompt_detail = custom_prompt if custom_prompt else self.prompt_detail
        return super().generate_topic_detail(self.prompt_sys, prompt_detail, text, topic)

this is tariff agent 

i need to know if we upload more than one transcript how things are genrated the sumary ,the topics are we sending all file data at once or what