from agents.base_agent import BaseAgent

class OverallAgent(BaseAgent):
    prompt_sys = "Macy's Executive Business Analyst works for only provided documents"
    
    # Universal formatting enforcement line
    formatting_directive = """STRICTLY ENFORCE PLAIN TEXT FORMATTING:
2. **No Combined Elements**: Always space-separate numbers/words/units (3.3 billion not 3.3billion)
3. **Time Formatting**: Q4 2024 (never Q42024 or 4th quarter)
4. **Number Rules**: 
   - 3.3B → 3.3 billion
   - 150.2M → 150.2 million
   - 0.33 → keep exact decimal
5. **Text Integrity**: 
   - "net income" never "netincome"
   - "year-over-year" never "YOY" or "yearoveryear"
6. **Style Prohibition**: Do not make any word in Italic or any other format always make it in one format
7. **Bad Example**: 3.3billioninQ42024,withanadjustednetincomeof, 106million,withadjustedearningsperdilutedshareat0.95 → REJECTED
8. **ALLWAYS Required Format**: 3.3 billion in Q4 2024, with an adjusted net income of → ACCEPTED


APPLY THESE RULES TO ALL NUMERICAL MENTIONS AND SENTENCE STRUCTURES."""
    
    prompt_analyze = """You are a senior business analyst working directly for Macy's CEO. Analyze the provided earnings call transcript and identify all significant business topics across all operational areas.

Identify key topics including but not limited to:
- Financial performance metrics
- Consumer behavior trends
- Supply chain and logistics
- Tariff and trade impacts
- Technology investments
- Strategic initiatives
- Competitive landscape

For each topic, provide:
* A clear topic name (e.g., "E-commerce Growth", "Supply Chain Digitization")
* Concise 100-word summary providing key numbers and operational impacts
* Relevance score (1-10) based on strategic importance to Macy's
* Strictly preserve original numerical and textual formatting: maintain exact decimals (0.33), required spacing between numbers/text (0.11 → '0.11'), and verbatim terminology without concatenation or rephrasing.
* For every topic, ALWAYS provide excerpts mentioned under the heading "Ground Truth" as follows:
        - For each document where the topic appears, extract the exact, unmodified sentence(s) or paragraph(s) from the original document. 
        - Format each excerpt as: **Ground Truth:** "Exact text from the document..." (Document Name)
        - If a topic is present in multiple documents, include a separate "Ground Truth" excerpt for each document.
        - Do NOT paraphrase, summarize, or alter the original text in any way. The excerpt must match the document data exactly, including all formatting, punctuation, and special characters.
        - If the LLM cannot find an exact match in the document, do NOT fabricate or modify the excerpt—only include what is present in the documents.

Format each topic as:
- Topic Name: Summary [Relevance Score]

Example:
- Holiday Sales Performance: Q4 comparable sales up 4.2% driven by 18% digital growth [8]
- Inventory Optimization: Implementing AI demand forecasting to reduce carrying costs by $150M [7]

Document Text:""" + formatting_directive
    

    prompt_summary = """You are a senior analyst preparing a consolidated report for Macy's CEO from all documents provided.

Identify and describe all impactful business topics from all the provided documents, focusing on those with broad implications for Macy's operations, strategy, and financial performance. Combine and synthesize topics that appear in multiple documents.

1. **Overview**: Identify and describe impactful topics from all documents, especially those related to government policy or other external factors that could affect Macy's business operations.

2. For each topic:
    * Provide a concise summary by combining information from all relevant documents.
    * Include relevant examples, numerical data, or timelines from the documents.
    * Explain operational and strategic impacts—positive, negative, or neutral—specifically tailored to Macy's retail business.
    * Provide 3-5 actionable strategic recommendations for Macy's leadership.
    * For every topic, ALWAYS provide each excerpts mentioned under the heading "Ground Truth" as follows:
        - For each document where the topic appears, extract the exact, unmodified sentence(s) or paragraph(s) from the original document. 
        - Format each excerpt as: **Ground Truth:** "Exact text from the document..." (Document Name)
        - If a same topic is present in multiple documents, always include a separate "Ground Truth" excerpt for each document with proper line breaks.
        - Do NOT paraphrase, summarize, or alter the original text in any way. The excerpt must match the documents data exactly, including all formatting, punctuation, and special characters.
        - If the LLM cannot find an exact match in the document, do NOT fabricate or modify the excerpt—only include what is present in the document.

3. Format in clear with sections in black text with proper headings:
    * Executive Summary
    * Key Themes
    * Strategic Recommendations
    * Risk Assessment

Use the following Documents Analysis Data to combine similar topics from different documents and provide a comprehensive summary:
{topics_block}

Provide comprehensive yet concise analysis suitable for C-level decision making.""" + formatting_directive 
    
    prompt_detail ="""You are preparing an executive briefing on '{topic}' for Macy's leadership. Provide:

1. **Overview**: Context and importance of this topic across all documents
2. **Cross-Document Insights**: Always Provide Key findings from all relevant documents, Provide same text do not make any changes.
3. **Financial Impact**: Any disclosed numbers or estimates from all sources
4. **Operational Effects**: Combined impacts on retail ops, supply chain, and digital
5. **Recommendations**: 3-5 actionable strategies with timelines
6. For every topic, ALWAYS provide excerpts mentioned under the heading "Ground Truth" as follows:
        - For each document where the topic appears, extract the exact, unmodified sentence(s) or paragraph(s) from the original document. 
        - Format each excerpt as: **Ground Truth:** "Exact text from the document..." (Document Name)
        - If a topic is present in multiple documents, include a separate "Ground Truth" excerpt for each document.
        - Do NOT paraphrase, summarize, or alter the original text in any way. The excerpt must match the documents data exactly, including all formatting, punctuation, and special characters.
        - If the LLM cannot find an exact match in the document, do NOT fabricate or modify the excerpt—only include what is present in the documents.

Include specific examples from documents where available, noting which document they came from.
Strictly preserve original numerical and textual formatting: maintain exact decimals (0.33), required spacing between numbers/text (0.11 → '0.11'), and verbatim terminology without concatenation or rephrasing.

Relevant document excerpts:
{text}

Deliver response in structured with proper paragraphs and bold topics and headings.""" + formatting_directive

    def analyze_document(self, text, custom_prompt=None):
        print("Analyzing document with custom prompt:", custom_prompt)
        prompt_analyze = custom_prompt if custom_prompt else self.prompt_analyze
        return super().analyze_document(self.prompt_sys, prompt_analyze, text)

    def generate_consolidated_summary_from_topics(self, processed_files, custom_prompt=None):
        print("Generating consolidated summary with custom prompt:", custom_prompt)
        prompt_summary = custom_prompt if custom_prompt else self.prompt_summary
        return super().generate_consolidated_summary_from_topics(self.prompt_sys, prompt_summary, processed_files)

    def generate_topic_detail(self, text, topic,custom_prompt=None):
        print("Generating topic detail with custom prompt:", custom_prompt)
        prompt_detail = custom_prompt if custom_prompt else self.prompt_detail
        return super().generate_topic_detail(self.prompt_sys, prompt_detail, text, topic)


this is overall agent which ahve custome prompts and normal prompts

from utils.pdf_utils import parse_topics

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

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

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

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

        prompt_user = prompt_user_template.format(topics_block=topics_block)

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

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


this is the base 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


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,
            'active_topic': None,
            'chat_history': [],
            'selected_files': []
        }
    qa_state = st.session_state.qa_page

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

    # Initialize 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 structured Q&A content
                                qa_data = st.session_state.agent_qa.analyze_question(text)
                                
                                # Store the structured data
                                qa_state['processed_files'][filename] = {
                                    'text': text,
                                    'qa_data': qa_data
                                }
                                st.toast(f"✅ Processed {filename}")
                                
                                # Automatically select the first topic
                                if not qa_state.get('error') and qa_data.get('topics'):
                                    qa_state['active_topic'] = qa_data['topics'][0]
                            else:
                                st.error(f"❌ Failed to extract text for {filename}")
                        except Exception as e:
                            st.error(f"Error processing {filename}: {str(e)}")

            # 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
            # Automatically select the first topic for this document
            file_data = qa_state['processed_files'][filename]
            qa_data = file_data['qa_data']
            if not qa_data.get('error') and qa_data.get('topics'):
                qa_state['active_topic'] = qa_data['topics'][0]

    # --- Topic Sidebar ---
    if qa_state['active_file']:
        file_data = qa_state['processed_files'][qa_state['active_file']]
        qa_data = file_data['qa_data']
        
        st.sidebar.header("📌 Topics")
        
        # Handle error case
        if "error" in qa_data:
            st.sidebar.error("Error in document processing")
            
        elif 'topics' in qa_data:
            topics = qa_data['topics']
            
            # Display each topic in the sidebar
            for topic in topics:
                btn_label = f"{topic['topic']} ({topic['mention_count']} mentions)"
                if st.sidebar.button(
                    btn_label,
                    key=f"topic_{topic['topic']}_{qa_state['active_file']}",
                    use_container_width=True,
                ):
                    qa_state['active_topic'] = topic
        else:
            st.sidebar.info("No topics found in this document")

    # --- Q&A Display ---
    st.subheader("📝 Extracted Q&A Content")
    
    if qa_state['active_file']:
        file_data = qa_state['processed_files'][qa_state['active_file']]
        qa_data = file_data['qa_data']
        st.markdown(f"**Document:** {qa_state['active_file']}")
        
        # Handle error case
        if "error" in qa_data:
            st.error(f"Error processing document: {qa_data['error']}")
            if "raw_response" in qa_data:
                st.text_area("Raw AI Response", qa_data["raw_response"], height=300)
        
        # Display key takeaways
        elif 'key_takeaways' in qa_data and qa_data['key_takeaways']:
            st.subheader("Key Takeaways")
            for takeaway in qa_data['key_takeaways']:
                st.markdown(f"- {takeaway}")
        
        # Display selected topic
        if qa_state['active_topic']:
            topic = qa_state['active_topic']
            st.divider()
            st.subheader(topic['topic'])
            st.markdown(f"**Mentioned:** {topic['mention_count']} times")
            st.divider()
            
            # Display questions for this topic
            if 'questions' in topic and topic['questions']:
                for idx, question in enumerate(topic['questions']):
                    firm = question.get('firm', 'Unknown Firm')
                    st.markdown(f"**Firm:** {firm}")
                    st.markdown(f"**Question {idx+1}:** {question['question']}")
                    st.write("---")
            else:
                st.info("No questions found for this topic")
        elif 'topics' in qa_data and qa_data['topics']:
            # If no topic selected but topics exist, show the first one
            qa_state['active_topic'] = qa_data['topics'][0]
            st.experimental_rerun()
        else:
            st.info("No Q&A content found for this document")
    else:
        st.info("Upload documents to see extracted Q&A content")

    st.divider()

    # --- 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))],
        key="qa_chat_selector"
    )
    qa_state['selected_files'] = selected_files

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

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

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

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

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

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

                response_placeholder.markdown(full_response)

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

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

    # 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,
            'active_topic': None,
            'chat_history': [],
            'selected_files': []
        })
        st.experimental_rerun()


if __name__ == "__main__":
    main()

this is main.py

my custom prompt for all three analyze_document, generate_consolidated_summary_from_topics, and generate_topic_detail are not working properly in the frontend i can change all three and ask to genrate topics with custom prompt but it is not working 

correct my code and give me correct full code