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":
                        topics = agent.analyze_document(text, custom_prompt)
                    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("❓ Extract 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]

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

if __name__ == "__main__":
    main()