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


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)

        # 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"]:
            if "agent_overall" not in st.session_state:
                is_azure = st.session_state.provider == "Azure GPT"
                st.session_state.agent_overall = OverallAgent(st.session_state.client, is_azure)
            if "agent_tariff" not in st.session_state:
                is_azure = st.session_state.provider == "Azure GPT"
                st.session_state.agent_tariff = TariffTradeAgent(st.session_state.client, is_azure)
            if "agent_monthly_performance" not in st.session_state:
                is_azure = st.session_state.provider == "Azure GPT"
                st.session_state.agent_monthly_performance = MonthlyPerformanceAgent(st.session_state.client, is_azure)
            if "agent_consumer_health" not in st.session_state:
                is_azure = st.session_state.provider == "Azure GPT"
                st.session_state.agent_consumer_health = ConsumerHealthAgent(st.session_state.client, is_azure)
            if "chat_agent" not in st.session_state:
                is_azure = st.session_state.provider == "Azure GPT"
                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

    sess['last_analysis_mode'] = mode
    sess['active_file'] = None
    sess['active_topic'] = None
    st.success(f"Analysis and summary generated for {mode.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': {},
            '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,
            '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
            }
        }

    tab_overall, tab_tariff, tab_monthly_performance, tab_consumer_health = st.tabs([
        "📑 Overall", "🔍 Tariff", "📈 Monthly Performance", "🩺 Consumer Health"
    ])

    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"

    st.subheader("Quick Analysis")
    col0, col1, col2, col3 = st.columns(4)
    analyze_overall = col0.button("📑 Overall Summary")
    analyze_tariff = col1.button("🔍 Analyze Tariffs")
    analyze_monthly_performance = col2.button("📈 Analyze Monthly Performance")
    analyze_consumer_health = col3.button("🩺 Analyze Consumer Health")

    if generate_mode or analyze_overall or analyze_tariff or analyze_monthly_performance or analyze_consumer_health:
        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 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

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

        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('consumer_health','consumer health').replace('monthly_performance','monthly performance')} analysis")
        else:
            if mode == "overall":
                agent.prompt_summary = OverallAgent.prompt_summary
                agent.prompt_analyze = OverallAgent.prompt_analyze
                agent.prompt_detail = OverallAgent.prompt_detail
            elif mode == "tariff":
                agent.prompt_summary = TariffTradeAgent.prompt_summary
                agent.prompt_analyze = TariffTradeAgent.prompt_analyze
                agent.prompt_detail = TariffTradeAgent.prompt_detail
            elif mode == "monthly_performance":
                agent.prompt_summary = MonthlyPerformanceAgent.prompt_summary
                agent.prompt_analyze = MonthlyPerformanceAgent.prompt_analyze
                agent.prompt_detail = MonthlyPerformanceAgent.prompt_detail
            elif mode == "consumer_health":
                agent.prompt_summary = ConsumerHealthAgent.prompt_summary
                agent.prompt_analyze = ConsumerHealthAgent.prompt_analyze
                agent.prompt_detail = ConsumerHealthAgent.prompt_detail

        with st.spinner(f"🔍 Analyzing {mode.replace('consumer_health', 'consumer health').replace('monthly_performance', 'monthly performance')} 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
        }
    }

    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:  # tariff, monthly_performance, consumer_health
        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"
        ]
        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': {},
            '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,
            'last_analysis_mode': None,
            'custom_prompt_enabled': {},
            'custom_prompt_text': {}
        }

        st.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.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.rerun()


if __name__ == "__main__":
    main()

with this code as we have differenet agent like tariff and Analyze Monthly Performance, Analyze Consumer Health can you add for me these below agents also import them properly and add them properly where needed 

QTD Performance
Category Performance
Competitive Landscape
Weather
FY25 Expectations
Geopolitical Uncertainty
Wholesale