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.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


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)

        # *** 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)

        # 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_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 == "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_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_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=["pdf"],
        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
            },
            '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_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", "🔍 Tariff", "📈 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("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_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 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")
    

    if generate_mode or analyze_overall or analyze_tariff 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
            "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 == "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,
                "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': "Tariff Topics",
            'main_header': "🌟 Tariff Analysis",
            'agent': st.session_state.agent_tariff
        },
        '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_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_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_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()

can you make this code in different different files seprate them as per working as well as provide the what the function is doing explanation on it so this become easy for other developer