from agents.base_agent import BaseAgent

class QAAgent(BaseAgent):
    prompt_sys = "Expert in extracting all Questions and Answers from transcripts for Macy's CEO"
  

    prompt_analyze = """
You are an expert in analyzing transcripts of earnings calls or business meetings. Your task is to carefully review the provided transcript and extract EVERY question and its corresponding answer EXACTLY as they appear, including the speaker information.

Your output must follow this 3-part format:

---

**Key Takeaways for Management:**  
Summarize analyst concerns, repeated themes, or signals from their questions (3–5 bullets max).  
Each bullet should reflect the tone, focus, or concern based on the questions.

---

**Top Topics**  
Group all questions by inferred topic. For each topic - Clear sections using ## Topic:
- Provide the topic title in **Topic:**.
- If the topic is unclear, use a placeholder like "General Questions".
- Mention how many times it was brought up (e.g., Mentioned 3 times).
- Ensure topics are **sorted in descending order** based on how frequently they were mentioned.
- List each question under it using this format and leave a line break between questions:
  Topic:Title of the Topic
  Mentioned [number] times
  Firm Name: The Firm Name (if available)
            **Question 1:** [exact question text]  
  
  Firm Name: The Firm Name (if available)
            **Question 2:** [exact question text] 
            
            
         
and so on for all questions under that topic.
--------------------------------         
            
  Topic:Title of the Topic
  Mentioned [number] times
  Firm Name: The Firm Name (if available)
            **Question 1:** [exact question text]  
            
  Firm Name: The Firm Name (if available)
            **Question 2:** [exact question text]  
            

         
and so on for all questions under that topic.
------------------                    


Additional Instructions:
- Do not paraphrase or summarize Q&A content – use exact text.
- Group based on the logical themes of the questions.
- Do not skip any questions, including follow-ups.
- Keep numerical data, names, and specific terminology exactly as stated.
- Questions without clear answers should still be included and marked as such.

Transcript:
{text}
"""



    def analyze_question(self, text, custom_prompt=None):
        prompt_to_use = custom_prompt if custom_prompt else self.prompt_analyze
        return super().analyze_question(
            self.prompt_sys, 
            prompt_to_use, 
            text
        )


thi is my qa agent 

# import streamlit as st
# from io import BytesIO
# from datetime import datetime

# # 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.sales_agent import SalesAgent
# from agents.supply_chain_agent import SupplyChainAgent
# from agents.overall_agent import OverallAgent
# from agents.enhanced_chat_agent import EnhancedChatAgent
# from agents.qa_agent import QAAgent  # Add QA agent import

# 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:
#                     if mode == "overall":
#                         # FIXED: Assign processed_files for overall mode
#                         topics = agent.analyze_document(text, custom_prompt)
#                         processed_files[filename] = {
#                             "text": text,
#                             "topics": topics,
#                             "context": text,
#                         }
#                     elif mode == "qa":
#                         # QA mode uses analyze_question instead
#                         qa_content = agent.analyze_question(text, custom_prompt)
#                         processed_files[filename] = {
#                             "text": text,
#                             "qa": qa_content,  # Store QA content directly
#                         }
#                     else:
#                         topics = agent.analyze_document(text, custom_prompt)
#                         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 (skip for QA)
#     if mode != "qa":
#         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 == "sales":
#         sess['consolidated_summary_sales'] = summary
#     elif mode == "supplychain":
#         sess['consolidated_summary_supplychain'] = summary

#     sess['last_analysis_mode'] = mode
#     sess['active_file'] = None
#     sess['active_topic'] = None
#     st.success(f"Analysis and summary generated for {mode.replace('supplychain', 'supply chain')}.")

# def main():
#     st.set_page_config(page_title="Macy's Document Analysis Pro", layout="wide")
#     st.title("📄 Macy's Document Analysis & Chat Tool")

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

#     # Session Initialization
#     if "session" not in st.session_state:
#         st.session_state.session = {
#             'uploaded_files': {},
#             'processed_files_overall': {},
#             'processed_files_tariff': {},
#             'processed_files_sales': {},
#             'processed_files_supplychain': {},
#             'processed_files_qa': {},  # New QA processed files
#             'active_file': None,
#             'active_topic': None,
#             'topic_details_cache': {},
#             'chat_history': [],
#             'consolidated_summary_overall': None,
#             'consolidated_summary_tariff': None,
#             'consolidated_summary_sales': None,
#             'consolidated_summary_supplychain': None,
#             'last_analysis_mode': None,
#             'custom_prompt_enabled': {},
#             'custom_prompt_text': {}
#         }

#     if "client" not in st.session_state or st.session_state.get("provider") != provider_select:
#         if provider_select == "Azure GPT":
#             st.session_state.client = AzureClient()
#             is_azure = True
#         else:
#             st.session_state.client = GeminiClient()
#             is_azure = False

#         # Initialize agents (add QA agent)
#         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_sales = SalesAgent(st.session_state.client, is_azure)
#         st.session_state.agent_supplychain = SupplyChainAgent(st.session_state.client, is_azure)
#         st.session_state.agent_qa = QAAgent(st.session_state.client, is_azure)  # New QA agent
#         st.session_state.chat_agent = EnhancedChatAgent(st.session_state.client, is_azure)
#         st.session_state.provider = provider_select

#     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..."):
#                 # Always run initial overall analysis for new files
#                 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")
    
#     # Initialize custom prompts in session state
#     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
#             },
#             'sales': {
#                 'summary': SalesAgent.prompt_summary,
#                 'analyze': SalesAgent.prompt_analyze,
#                 'detail': SalesAgent.prompt_detail
#             },
#             'supplychain': {
#                 'summary': SupplyChainAgent.prompt_summary,
#                 'analyze': SupplyChainAgent.prompt_analyze,
#                 'detail': SupplyChainAgent.prompt_detail
#             },
#             'qa': {  # New QA prompts
#                 'analyze': QAAgent.prompt_analyze
#             }
#         }
    
#     # Create tabs for each agent's custom prompts (add QA tab)
#     tab_overall, tab_tariff, tab_sales, tab_supply, tab_qa = st.tabs([
#         "📑 Overall", "🔍 Tariff", "📈 Sales", "🚚 Supply Chain", "❓ Q&A"
#     ])
    
#     # Store which mode we're generating for
#     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_sales:
#         st.subheader("Sales Prompts")
#         use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_sales")
        
#         if use_custom:
#             sess['custom_prompts']['sales']['summary'] = st.text_area(
#                 "Summary Prompt:",
#                 value=sess['custom_prompts']['sales']['summary'],
#                 height=150,
#                 key="sales_prompt_summary"
#             )
#             sess['custom_prompts']['sales']['analyze'] = st.text_area(
#                 "Analysis Prompt:",
#                 value=sess['custom_prompts']['sales']['analyze'],
#                 height=150,
#                 key="sales_prompt_analyze"
#             )
#             sess['custom_prompts']['sales']['detail'] = st.text_area(
#                 "Detail Prompt:",
#                 value=sess['custom_prompts']['sales']['detail'],
#                 height=150,
#                 key="sales_prompt_detail"
#             )
            
#             if st.button("🔄 Generate with Custom Prompts", key="generate_sales"):
#                 generate_mode = "sales"
    
#     with tab_supply:
#         st.subheader("Supply Chain Prompts")
#         use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_supply")
        
#         if use_custom:
#             sess['custom_prompts']['supplychain']['summary'] = st.text_area(
#                 "Summary Prompt:",
#                 value=sess['custom_prompts']['supplychain']['summary'],
#                 height=150,
#                 key="supply_prompt_summary"
#             )
#             sess['custom_prompts']['supplychain']['analyze'] = st.text_area(
#                 "Analysis Prompt:",
#                 value=sess['custom_prompts']['supplychain']['analyze'],
#                 height=150,
#                 key="supply_prompt_analyze"
#             )
#             sess['custom_prompts']['supplychain']['detail'] = st.text_area(
#                 "Detail Prompt:",
#                 value=sess['custom_prompts']['supplychain']['detail'],
#                 height=150,
#                 key="supply_prompt_detail"
#             )
            
#             if st.button("🔄 Generate with Custom Prompts", key="generate_supplychain"):
#                 generate_mode = "supplychain"
                
#     # New QA tab
#     with tab_qa:
#         st.subheader("Q&A Extraction Prompts")
#         use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_qa")
        
#         if use_custom:
#             sess['custom_prompts']['qa']['analyze'] = st.text_area(
#                 "Q&A Extraction Prompt:",
#                 value=sess['custom_prompts']['qa']['analyze'],
#                 height=300,
#                 key="qa_prompt_analyze"
#             )
            
#             if st.button("🔄 Generate with Custom Prompts", key="generate_qa"):
#                 generate_mode = "qa"

#     # --- Analysis Buttons ---
#     st.subheader("Quick Analysis")
#     col0, col1, col2, col3, col4 = st.columns(5)  # Add extra column for QA
#     analyze_overall = col0.button("📑 Overall Summary")
#     analyze_tariff = col1.button("🔍 Analyze Tariffs")
#     analyze_sales = col2.button("📈 Analyze Sales")
#     analyze_supplychain = col3.button("🚚 Supply Chain")
#     analyze_qa = col4.button("❓ Q & A")  # New QA button

#     # Handle generation requests
#     if generate_mode or analyze_overall or analyze_tariff or analyze_sales or analyze_supplychain or analyze_qa:
#         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
#             "sales" if analyze_sales else
#             "supplychain" if analyze_supplychain else
#             "qa" if analyze_qa else None
#         )
        
#         agent = None
#         processed_files = None

#         if mode == "overall":
#             agent = st.session_state.agent_overall
#             processed_files = sess['processed_files_overall']
#         elif mode == "tariff":
#             agent = st.session_state.agent_tariff
#             processed_files = sess['processed_files_tariff']
#         elif mode == "sales":
#             agent = st.session_state.agent_sales
#             processed_files = sess['processed_files_sales']
#         elif mode == "supplychain":
#             agent = st.session_state.agent_supplychain
#             processed_files = sess['processed_files_supplychain']
#         elif mode == "qa":  # New QA mode
#             agent = st.session_state.agent_qa
#             processed_files = sess['processed_files_qa']

#         # Update agent with custom prompts if used
#         if generate_mode:
#             if mode != "qa":
#                 agent.prompt_summary = sess['custom_prompts'][mode]['summary']
#                 agent.prompt_analyze = sess['custom_prompts'][mode]['analyze']
#                 agent.prompt_detail = sess['custom_prompts'][mode]['detail']
#             else:
#                 agent.prompt_analyze = sess['custom_prompts'][mode]['analyze']
#             st.toast(f"Using custom prompts for {mode} analysis")
#         else:
#             # Reset to default prompts
#             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 == "sales":
#                 agent.prompt_summary = SalesAgent.prompt_summary
#                 agent.prompt_analyze = SalesAgent.prompt_analyze
#                 agent.prompt_detail = SalesAgent.prompt_detail
#             elif mode == "supplychain":
#                 agent.prompt_summary = SupplyChainAgent.prompt_summary
#                 agent.prompt_analyze = SupplyChainAgent.prompt_analyze
#                 agent.prompt_detail = SupplyChainAgent.prompt_detail
#             elif mode == "qa":
#                 agent.prompt_analyze = QAAgent.prompt_analyze

#         with st.spinner(f"🔍 Analyzing {mode.replace('supplychain', 'supply chain').replace('qa', 'Q&A')} documents..."):
#             analyze_documents(
#                 mode, 
#                 agent, 
#                 processed_files, 
#                 sess
#             )

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

#     # Consolidated data setup (add QA mode)
#     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
#         },
#         'sales': {
#             'files': sess['processed_files_sales'],
#             'summary': sess['consolidated_summary_sales'],
#             'sidebar_header': "Sales Topics",
#             'main_header': "🌟 Sales Analysis",
#             'agent': st.session_state.agent_sales
#         },
#         'supplychain': {
#             'files': sess['processed_files_supplychain'],
#             'summary': sess['consolidated_summary_supplychain'],
#             'sidebar_header': "Supply Chain Topics",
#             'main_header': "🌟 Supply Chain Analysis",
#             'agent': st.session_state.agent_supplychain
#         },
#         'qa': {  # New QA config
#             'files': sess['processed_files_qa'],
#             'summary': None,  # QA doesn't have consolidated summary
#             'sidebar_header': "Documents with Q&A",
#             'main_header': "❓ Extracted Q&A",
#             'agent': st.session_state.agent_qa
#         }
#     }

#     current_mode = mode_config[last_mode]

#     # === Fix: Auto select first QA file if none selected ===
#     if last_mode == "qa":
#         if sess['active_file'] is None and current_mode['files']:
#             first_file = next(iter(current_mode['files']))
#             sess['active_file'] = first_file

#     # --- 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']
#                 })

#         a = 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 = a[:3] , a[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 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}", # Changed key to avoid collision
#                     use_container_width=True,
#                 ):
#                     sess['active_topic'] = topic_name
#                     sess['active_file'] = None
                    
#     elif last_mode == "qa":  # New QA sidebar
#         for filename in current_mode['files']:
#             if st.sidebar.button(
#                 f"📄 {filename}",
#                 key=f"qa_file_{filename}",
#                 use_container_width=True,
#             ):
#                 sess['active_file'] = filename
#                 sess['active_topic'] = None

#     else: # For Tariff, Sales, Supply Chain modes
#         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']}", # Changed key
#                             type=style,
#                             use_container_width=True,
#                         ):
#                             sess['active_file'] = filename
#                             sess['active_topic'] = topic['name']

#     # --- Topic 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
#                 except Exception as e:
#                     st.error(f"Analysis error: {str(e)}")
                    
#     elif last_mode == "qa" and sess['active_file']:  # New QA display
#         filename = sess['active_file']
#         qa_content = current_mode['files'][filename].get('qa', '')
#         if qa_content:
#             st.markdown(qa_content)
#         else:
#             st.info("No Q&A content found for this document")
            
#     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
#             except Exception as e:
#                 st.error(f"Analysis error: {str(e)}")
#     else:
#         st.info("Select a topic from the sidebar")

#     st.markdown("---")

#     # --- Consolidated Summary ---
#     # Skip for QA mode
#     if last_mode != "qa":
#         st.header(current_mode['main_header'])
#         if current_mode['summary']:
#             st.markdown(current_mode['summary'])
#         else:
#             st.info("Summary not available yet")

#     st.divider()

#     # --- Chat Interface ---
#     st.subheader("💬 Document Chat")
    
#     # File selection
#     available_files = list(current_mode['files'].keys())
#     selected_files = st.multiselect(
#         "Select documents for context:",
#         available_files,
#         key=f"chat_selector_{last_mode}"
#     )

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

#     # Question handling
#     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:
#             # Fix: For QA mode, use 'text' instead of 'context' for chat context in files
#             if last_mode == "qa":
#                 context = "\n\n".join([
#                     f"=== {fname} ===\n{current_mode['files'][fname].get('text', '')}"
#                     for fname in selected_files
#                 ])
#             else:
#                 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
                
#                 # Ensure custom_prompt_text_for_mode is available here if needed for chat_agent
#                 # For now, it's used only for analysis in analyze_documents, adjust if chat_agent needs it.
#                 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(),
#             })

#     # Reset button
#     if st.button("🔄 Reset Session", type="secondary"):
#         # Preserve long-lived objects like client and agents to avoid overhead on re-init
#         preserved_keys = [
#             "client", "agent_overall", "agent_tariff",
#             "agent_sales", "agent_supplychain", "agent_qa", "chat_agent", "provider"
#         ]
#         preserved = {k: st.session_state[k] for k in preserved_keys if k in st.session_state}
        
#         # Clear entire session state
#         st.session_state.clear()
        
#         # Restore preserved keys back
#         st.session_state.update(preserved)

#         # Reset session-specific stored data completely
#         st.session_state.session = {
#             'uploaded_files': {},
#             'processed_files_overall': {},
#             'processed_files_tariff': {},
#             'processed_files_sales': {},
#             'processed_files_supplychain': {},
#             'processed_files_qa': {},  # New QA processed files
#             'active_file': None,
#             'active_topic': None,
#             'topic_details_cache': {},
#             'chat_history': [],
#             'consolidated_summary_overall': None,
#             'consolidated_summary_tariff': None,
#             'consolidated_summary_sales': None,
#             'consolidated_summary_supplychain': None,
#             'last_analysis_mode': None,
#             'custom_prompt_enabled': {},
#             'custom_prompt_text': {}
#         }
        
#         # Trigger a rerun for fresh UI without old content
#         st.experimental_rerun()

# if __name__ == "__main__":
#     main()



# import streamlit as st
# from io import BytesIO
# from datetime import datetime

# # 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.sales_agent import SalesAgent
# from agents.supply_chain_agent import SupplyChainAgent
# from agents.overall_agent import OverallAgent
# from agents.enhanced_chat_agent import EnhancedChatAgent
# from agents.qa_agent import QAAgent

# 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:
#                     topics = agent.analyze_document(text, custom_prompt)
#                     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 == "sales":
#         sess['consolidated_summary_sales'] = summary
#     elif mode == "supplychain":
#         sess['consolidated_summary_supplychain'] = summary

#     sess['last_analysis_mode'] = mode
#     sess['active_file'] = None
#     sess['active_topic'] = None
#     st.success(f"Analysis and summary generated for {mode.replace('supplychain', 'supply chain')}.")

# def main():
#     st.set_page_config(page_title="Macy's Document Analysis Pro", layout="wide")
#     st.title("📊 Macy's Document Analysis Dashboard")
    
#     # Navigation
#     st.sidebar.title("Navigation")
#     page = st.sidebar.radio("Go to", ["Analysis Dashboard", "Q&A Assistant"])
    
#     if page == "Analysis Dashboard":
#         dashboard_page()
#     elif page == "Q&A Assistant":
#         qa_page()

# def dashboard_page():
#     # Mode selection
#     provider_select = st.radio("Select AI Provider", ["Azure GPT", "Google Gemini"])

#     # Session Initialization
#     if "session" not in st.session_state:
#         st.session_state.session = {
#             'uploaded_files': {},
#             'processed_files_overall': {},
#             'processed_files_tariff': {},
#             'processed_files_sales': {},
#             'processed_files_supplychain': {},
#             'active_file': None,
#             'active_topic': None,
#             'topic_details_cache': {},
#             'chat_history': [],
#             'consolidated_summary_overall': None,
#             'consolidated_summary_tariff': None,
#             'consolidated_summary_sales': None,
#             'consolidated_summary_supplychain': None,
#             'last_analysis_mode': None,
#             'custom_prompt_enabled': {},
#             'custom_prompt_text': {}
#         }

#     if "client" not in st.session_state or st.session_state.get("provider") != provider_select:
#         if provider_select == "Azure GPT":
#             st.session_state.client = AzureClient()
#             is_azure = True
#         else:
#             st.session_state.client = GeminiClient()
#             is_azure = False

#         # Initialize 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_sales = SalesAgent(st.session_state.client, is_azure)
#         st.session_state.agent_supplychain = SupplyChainAgent(st.session_state.client, is_azure)
#         st.session_state.chat_agent = EnhancedChatAgent(st.session_state.client, is_azure)
#         st.session_state.provider = provider_select

#     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..."):
#                 # Always run initial overall analysis for new files
#                 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")
    
#     # Initialize custom prompts in session state
#     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
#             },
#             'sales': {
#                 'summary': SalesAgent.prompt_summary,
#                 'analyze': SalesAgent.prompt_analyze,
#                 'detail': SalesAgent.prompt_detail
#             },
#             'supplychain': {
#                 'summary': SupplyChainAgent.prompt_summary,
#                 'analyze': SupplyChainAgent.prompt_analyze,
#                 'detail': SupplyChainAgent.prompt_detail
#             }
#         }
    
#     # Create tabs for each agent's custom prompts
#     tab_overall, tab_tariff, tab_sales, tab_supply = st.tabs([
#         "📑 Overall", "🔍 Tariff", "📈 Sales", "🚚 Supply Chain"
#     ])
    
#     # Store which mode we're generating for
#     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_sales:
#         st.subheader("Sales Prompts")
#         use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_sales")
        
#         if use_custom:
#             sess['custom_prompts']['sales']['summary'] = st.text_area(
#                 "Summary Prompt:",
#                 value=sess['custom_prompts']['sales']['summary'],
#                 height=150,
#                 key="sales_prompt_summary"
#             )
#             sess['custom_prompts']['sales']['analyze'] = st.text_area(
#                 "Analysis Prompt:",
#                 value=sess['custom_prompts']['sales']['analyze'],
#                 height=150,
#                 key="sales_prompt_analyze"
#             )
#             sess['custom_prompts']['sales']['detail'] = st.text_area(
#                 "Detail Prompt:",
#                 value=sess['custom_prompts']['sales']['detail'],
#                 height=150,
#                 key="sales_prompt_detail"
#             )
            
#             if st.button("🔄 Generate with Custom Prompts", key="generate_sales"):
#                 generate_mode = "sales"
    
#     with tab_supply:
#         st.subheader("Supply Chain Prompts")
#         use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_supply")
        
#         if use_custom:
#             sess['custom_prompts']['supplychain']['summary'] = st.text_area(
#                 "Summary Prompt:",
#                 value=sess['custom_prompts']['supplychain']['summary'],
#                 height=150,
#                 key="supply_prompt_summary"
#             )
#             sess['custom_prompts']['supplychain']['analyze'] = st.text_area(
#                 "Analysis Prompt:",
#                 value=sess['custom_prompts']['supplychain']['analyze'],
#                 height=150,
#                 key="supply_prompt_analyze"
#             )
#             sess['custom_prompts']['supplychain']['detail'] = st.text_area(
#                 "Detail Prompt:",
#                 value=sess['custom_prompts']['supplychain']['detail'],
#                 height=150,
#                 key="supply_prompt_detail"
#             )
            
#             if st.button("🔄 Generate with Custom Prompts", key="generate_supplychain"):
#                 generate_mode = "supplychain"

#     # --- Analysis Buttons ---
#     st.subheader("Quick Analysis")
#     col0, col1, col2, col3 = st.columns(4)
#     analyze_overall = col0.button("📑 Overall Summary")
#     analyze_tariff = col1.button("🔍 Analyze Tariffs")
#     analyze_sales = col2.button("📈 Analyze Sales")
#     analyze_supplychain = col3.button("🚚 Supply Chain")

#     # Handle generation requests
#     if generate_mode or analyze_overall or analyze_tariff or analyze_sales or analyze_supplychain:
#         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
#             "sales" if analyze_sales else
#             "supplychain" if analyze_supplychain else None
#         )
        
#         agent = None
#         processed_files = None

#         if mode == "overall":
#             agent = st.session_state.agent_overall
#             processed_files = sess['processed_files_overall']
#         elif mode == "tariff":
#             agent = st.session_state.agent_tariff
#             processed_files = sess['processed_files_tariff']
#         elif mode == "sales":
#             agent = st.session_state.agent_sales
#             processed_files = sess['processed_files_sales']
#         elif mode == "supplychain":
#             agent = st.session_state.agent_supplychain
#             processed_files = sess['processed_files_supplychain']

#         # Update agent with custom prompts if used
#         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} analysis")
#         else:
#             # Reset to default prompts
#             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 == "sales":
#                 agent.prompt_summary = SalesAgent.prompt_summary
#                 agent.prompt_analyze = SalesAgent.prompt_analyze
#                 agent.prompt_detail = SalesAgent.prompt_detail
#             elif mode == "supplychain":
#                 agent.prompt_summary = SupplyChainAgent.prompt_summary
#                 agent.prompt_analyze = SupplyChainAgent.prompt_analyze
#                 agent.prompt_detail = SupplyChainAgent.prompt_detail

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

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

#     # Consolidated data setup
#     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
#         },
#         'sales': {
#             'files': sess['processed_files_sales'],
#             'summary': sess['consolidated_summary_sales'],
#             'sidebar_header': "Sales Topics",
#             'main_header': "🌟 Sales Analysis",
#             'agent': st.session_state.agent_sales
#         },
#         'supplychain': {
#             'files': sess['processed_files_supplychain'],
#             'summary': sess['consolidated_summary_supplychain'],
#             'sidebar_header': "Supply Chain Topics",
#             'main_header': "🌟 Supply Chain Analysis",
#             'agent': st.session_state.agent_supplychain
#         }
#     }

#     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: # For Tariff, Sales, Supply Chain modes
#         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 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
#                 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
#             except Exception as e:
#                 st.error(f"Analysis error: {str(e)}")
#     else:
#         st.info("Select a topic from the sidebar")

#     st.markdown("---")

#     # --- Consolidated Summary ---
#     st.header(current_mode['main_header'])
#     if current_mode['summary']:
#         st.markdown(current_mode['summary'])
#     else:
#         st.info("Summary not available yet")

#     st.divider()

#     # --- Chat Interface ---
#     st.subheader("💬 Document Chat")
    
#     # File selection
#     available_files = list(current_mode['files'].keys())
#     selected_files = st.multiselect(
#         "Select documents for context:",
#         available_files,
#         key=f"chat_selector_{last_mode}"
#     )

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

#     # Question handling
#     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(),
#             })

#     # Reset button
#     if st.button("🔄 Reset Session", type="secondary"):
#         # Preserve long-lived objects like client and agents to avoid overhead on re-init
#         preserved_keys = [
#             "client", "agent_overall", "agent_tariff",
#             "agent_sales", "agent_supplychain", "chat_agent", "provider"
#         ]
#         preserved = {k: st.session_state[k] for k in preserved_keys if k in st.session_state}
        
#         # Clear entire session state
#         st.session_state.clear()
        
#         # Restore preserved keys back
#         st.session_state.update(preserved)

#         # Reset session-specific stored data completely
#         st.session_state.session = {
#             'uploaded_files': {},
#             'processed_files_overall': {},
#             'processed_files_tariff': {},
#             'processed_files_sales': {},
#             'processed_files_supplychain': {},
#             'active_file': None,
#             'active_topic': None,
#             'topic_details_cache': {},
#             'chat_history': [],
#             'consolidated_summary_overall': None,
#             'consolidated_summary_tariff': None,
#             'consolidated_summary_sales': None,
#             'consolidated_summary_supplychain': None,
#             'last_analysis_mode': None,
#             'custom_prompt_enabled': {},
#             'custom_prompt_text': {}
#         }
        
#         # Trigger a rerun for fresh UI without old content
#         st.experimental_rerun()

# def qa_page():
#     st.title("❓ Macy's Q&A Assistant")
#     st.markdown("Upload documents to automatically extract Q&A content and chat with your documents.")
    
#     # Initialize session state for QA page
#     if 'qa_page' not in st.session_state:
#         st.session_state.qa_page = {
#             'uploaded_files': {},
#             'processed_files': {},
#             'active_file': None,
#             'chat_history': [],
#             'selected_files': []
#         }
#     qa_state = st.session_state.qa_page

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

#     # Initialize client if needed
#     if "client" not in st.session_state or st.session_state.get("provider") != provider_select:
#         if provider_select == "Azure GPT":
#             st.session_state.client = AzureClient()
#             is_azure = True
#         else:
#             st.session_state.client = GeminiClient()
#             is_azure = False
        
#         # Initialize agents
#         st.session_state.agent_qa = QAAgent(st.session_state.client, is_azure)
#         st.session_state.chat_agent = EnhancedChatAgent(st.session_state.client, is_azure)
#         st.session_state.provider = provider_select

#     # --- File Upload ---
#     with st.expander("📤 Upload Documents", expanded=True):
#         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:
#                                     # Use the QA Agent to generate Q&A content
#                                     qa_content = st.session_state.agent_qa.analyze_question(text)
#                                     qa_state['processed_files'][filename] = {
#                                         'text': text,
#                                         'qa': qa_content
#                                     }
#                                     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)}")
                
#                 # Auto-select the first file if none is selected
#                 if not qa_state['active_file'] and qa_state['processed_files']:
#                     qa_state['active_file'] = next(iter(qa_state['processed_files'].keys()))

#     # --- Document Selection ---
#     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

#     # --- Q&A Display ---
#     st.subheader("📝 Extracted Q&A Content")
#     if qa_state['active_file']:
#         file_data = qa_state['processed_files'][qa_state['active_file']]
#         st.markdown(f"**Document:** {qa_state['active_file']}")
#         st.markdown(file_data['qa'])
#     else:
#         st.info("Upload documents to see extracted Q&A content")

#     # --- Chat Interface ---
#     st.divider()
#     st.subheader("💬 Chat with Documents")
    
#     # File selection for chat context
#     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))],  # Auto-select first 3
#         key="qa_chat_selector"
#     )
#     qa_state['selected_files'] = selected_files

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

#     # Question handling
#     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:
#             # Prepare context from selected files
#             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 = ""
                
#                 # Generate response using the chat agent
#                 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(),
#             })

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

# if __name__ == "__main__":
#     main()

import streamlit as st
from io import BytesIO
from datetime import datetime

# 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.sales_agent import SalesAgent
from agents.supply_chain_agent import SupplyChainAgent
from agents.overall_agent import OverallAgent
from agents.enhanced_chat_agent import EnhancedChatAgent
from agents.qa_agent import QAAgent


def initialize_clients_and_agents(provider_select: str, mode: str = None):
    """
    Ensure 'client' and agents are initialized in st.session_state.
    If mode is 'dashboard', initialize dashboard agents.
    If mode is 'qa', initialize QA agent.
    If mode is None, initialize all.
    """

    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

        # For 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_sales = SalesAgent(st.session_state.client, is_azure)
        st.session_state.agent_supplychain = SupplyChainAgent(st.session_state.client, is_azure)

        # Chat agent (used on both pages)
        st.session_state.chat_agent = EnhancedChatAgent(st.session_state.client, is_azure)

        # QA agent only if mode is 'qa' or None (default to initialize all)
        if mode == "qa" or mode is None:
            st.session_state.agent_qa = QAAgent(st.session_state.client, is_azure)

    else:
        # If client/provider the same, still validate agents are present and initialize if missing
        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_sales" not in st.session_state:
                is_azure = st.session_state.provider == "Azure GPT"
                st.session_state.agent_sales = SalesAgent(st.session_state.client, is_azure)
            if "agent_supplychain" not in st.session_state:
                is_azure = st.session_state.provider == "Azure GPT"
                st.session_state.agent_supplychain = SupplyChainAgent(st.session_state.client, is_azure)
            if "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:
                    topics = agent.analyze_document(text, custom_prompt)
                    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 == "sales":
        sess['consolidated_summary_sales'] = summary
    elif mode == "supplychain":
        sess['consolidated_summary_supplychain'] = summary

    sess['last_analysis_mode'] = mode
    sess['active_file'] = None
    sess['active_topic'] = None
    st.success(f"Analysis and summary generated for {mode.replace('supplychain', 'supply chain')}.")


def main():
    st.set_page_config(page_title="Macy's Document Analysis Pro", layout="wide")
    
    # Navigation
    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 client and agents for dashboard
    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_sales': {},
            'processed_files_supplychain': {},
            'active_file': None,
            'active_topic': None,
            'topic_details_cache': {},
            'chat_history': [],
            'consolidated_summary_overall': None,
            'consolidated_summary_tariff': None,
            'consolidated_summary_sales': None,
            'consolidated_summary_supplychain': 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..."):
                # Always run initial overall analysis for new files
                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")

    # Initialize custom prompts in session state
    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
            },
            'sales': {
                'summary': SalesAgent.prompt_summary,
                'analyze': SalesAgent.prompt_analyze,
                'detail': SalesAgent.prompt_detail
            },
            'supplychain': {
                'summary': SupplyChainAgent.prompt_summary,
                'analyze': SupplyChainAgent.prompt_analyze,
                'detail': SupplyChainAgent.prompt_detail
            }
        }

    # Create tabs for each agent's custom prompts
    tab_overall, tab_tariff, tab_sales, tab_supply = st.tabs([
        "📑 Overall", "🔍 Tariff", "📈 Sales", "🚚 Supply Chain"
    ])

    # Store which mode we're generating for
    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_sales:
        st.subheader("Sales Prompts")
        use_custom = st.checkbox("✏️ Use Custom Prompts", key="custom_sales")

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

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

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

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

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

    # --- Analysis Buttons ---
    st.subheader("Quick Analysis")
    col0, col1, col2, col3 = st.columns(4)
    analyze_overall = col0.button("📑 Overall Summary")
    analyze_tariff = col1.button("🔍 Analyze Tariffs")
    analyze_sales = col2.button("📈 Analyze Sales")
    analyze_supplychain = col3.button("🚚 Supply Chain")

    # Handle generation requests
    if generate_mode or analyze_overall or analyze_tariff or analyze_sales or analyze_supplychain:
        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
            "sales" if analyze_sales else
            "supplychain" if analyze_supplychain else None
        )

        agent = None
        processed_files = None

        if mode == "overall":
            agent = st.session_state.agent_overall
            processed_files = sess['processed_files_overall']
        elif mode == "tariff":
            agent = st.session_state.agent_tariff
            processed_files = sess['processed_files_tariff']
        elif mode == "sales":
            agent = st.session_state.agent_sales
            processed_files = sess['processed_files_sales']
        elif mode == "supplychain":
            agent = st.session_state.agent_supplychain
            processed_files = sess['processed_files_supplychain']

        # Update agent with custom prompts if used
        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} analysis")
        else:
            # Reset to default prompts
            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 == "sales":
                agent.prompt_summary = SalesAgent.prompt_summary
                agent.prompt_analyze = SalesAgent.prompt_analyze
                agent.prompt_detail = SalesAgent.prompt_detail
            elif mode == "supplychain":
                agent.prompt_summary = SupplyChainAgent.prompt_summary
                agent.prompt_analyze = SupplyChainAgent.prompt_analyze
                agent.prompt_detail = SupplyChainAgent.prompt_detail

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

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

    # Consolidated data setup
    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
        },
        'sales': {
            'files': sess['processed_files_sales'],
            'summary': sess['consolidated_summary_sales'],
            'sidebar_header': "Sales Topics",
            'main_header': "🌟 Sales Analysis",
            'agent': st.session_state.agent_sales
        },
        'supplychain': {
            'files': sess['processed_files_supplychain'],
            'summary': sess['consolidated_summary_supplychain'],
            'sidebar_header': "Supply Chain Topics",
            'main_header': "🌟 Supply Chain Analysis",
            'agent': st.session_state.agent_supplychain
        }
    }

    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:  # For Tariff, Sales, Supply Chain modes
        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 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
                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
            except Exception as e:
                st.error(f"Analysis error: {str(e)}")
    else:
        st.info("Select a topic from the sidebar")

    st.markdown("---")

    # --- Consolidated Summary ---
    st.header(current_mode['main_header'])
    if current_mode['summary']:
        st.markdown(current_mode['summary'])
    else:
        st.info("Summary not available yet")

    st.divider()

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

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

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

    # Question handling
    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(),
            })

    # Reset button
    if st.button("🔄 Reset Session", type="secondary"):
        preserved_keys = [
            "client", "agent_overall", "agent_tariff",
            "agent_sales", "agent_supplychain", "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_sales': {},
            'processed_files_supplychain': {},
            'active_file': None,
            'active_topic': None,
            'topic_details_cache': {},
            'chat_history': [],
            'consolidated_summary_overall': None,
            'consolidated_summary_tariff': None,
            'consolidated_summary_sales': None,
            'consolidated_summary_supplychain': 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,
            'chat_history': [],
            'selected_files': []
        }
    qa_state = st.session_state.qa_page

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

    # Initialize client and QA agents
    initialize_clients_and_agents(provider_select, mode="qa")

    # --- File Upload ---
    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:
                                # Use the QA Agent to generate Q&A content
                                qa_content = st.session_state.agent_qa.analyze_question(text)
                                qa_state['processed_files'][filename] = {
                                    'text': text,
                                    'qa': qa_content
                                }
                                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)}")

            # Auto-select the first file if none is selected
            if not qa_state['active_file'] and qa_state['processed_files']:
                qa_state['active_file'] = next(iter(qa_state['processed_files'].keys()))

    # --- Document Selection ---
    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

    # --- Q&A Display ---
    st.subheader("📝 Extracted Q&A Content")
    if qa_state['active_file']:
        file_data = qa_state['processed_files'][qa_state['active_file']]
        st.markdown(f"**Document:** {qa_state['active_file']}")
        st.markdown(file_data['qa'])
    else:
        st.info("Upload documents to see extracted Q&A content")

    st.divider()

    # --- Chat Interface ---
    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))],  # Auto-select first 3 if available
        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 = ""

                # <-- Fix: unpack tuple from generate_answer() here -->
                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(),
            })

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


if __name__ == "__main__":
    main()

this is my main.py

now change in the qa part only where i need to get the topics and the mentioned time in the sidebar , all the topics for that we can change in the prompt and get the topics and the mentioned and there question seprately for each of them and then we can display the topics and mentioned time in sidebar on clicking of topics we cn display the question for it along with firm 
and for the first topic automoatically dispaly the question of it on screen 

give me full code  do not remove any other stuff