import streamlit as st
import os
import re
from datetime import datetime
from PyPDF2 import PdfReader
import openai
from io import BytesIO

# ===== Utility Functions =====
def extract_text(file):
    try:
        reader = PdfReader(file)
        return "\n".join([page.extract_text() or "" for page in reader.pages])
    except Exception as e:
        st.error(f"Error reading {file.name}: {str(e)}")
        return ""

# ===== Agent Classes =====
class TariffTradeAgent:
    def __init__(self, client, deployment_id):
        self.client = client
        self.deployment_id = deployment_id

    def analyze_document(self, text):
        prompt = """
You are a senior tariff and international trade policy analyst assisting a research team working for the CEO of Macy's. Identify all significant tariff and trade policy topics mentioned in the text that could impact Macy's business operations, supply chain, or market opportunities.

For each topic, provide:
- A clear topic name
- Concise 1-sentence summary describing its impact
- Relevance score (1-10)

Output format:
- Topic Name: Summary [Relevance Score]

Example:
- USMCA Trade Agreement: Expected tariff reductions improve apparel import costs [8]
- Chinese Tariffs Impact: Increased tariffs lead to higher electronic goods prices [7]

Text:
"""
        response = self.client.chat.completions.create(
            model=self.deployment_id,
            messages=[
                {"role": "system", "content": "Tariff analyst for Macy's CEO"},
                {"role": "user", "content": prompt + "\n\n" + text}
            ],
            temperature=0
        )
        return self._parse_topics(response.choices[0].message.content.strip())

    def _parse_topics(self, text):
        topics = []
        pattern = r'^- (.*?): (.*?) \[(\d+)\]$'
        for line in text.split('\n'):
            line = line.strip()
            match = re.match(pattern, line)
            if match:
                topics.append({
                    'name': match.group(1).strip(),
                    'summary': match.group(2).strip(),
                    'score': int(match.group(3).strip())
                })
        return topics

    def generate_consolidated_summary_from_topics(self, processed_files):
        # Build topic map and identify files without topics
        topic_map = {}
        all_files = set(processed_files.keys())
        files_with_topics = set()
        
        for filename, filedata in processed_files.items():
            for topic in filedata['topics']:
                name = topic['name']
                if name not in topic_map:
                    topic_map[name] = {
                        'files': set(),
                        'summaries': [],
                        'max_score': 0
                    }
                topic_map[name]['files'].add(filename)
                topic_map[name]['summaries'].append(topic['summary'])
                topic_map[name]['max_score'] = max(topic_map[name]['max_score'], topic['score'])
                files_with_topics.add(filename)

        # Prepare content note
        files_without_topics = all_files - files_with_topics
        content_note = (
            f"\n\n**Note:** The following documents contained content not captured in topics: {', '.join(files_without_topics)}. " 
            "Please review separately for additional insights." 
            if files_without_topics else ""
        )

        # Build topics text for prompt
        sorted_topics = sorted(topic_map.items(), key=lambda x: x[1]['max_score'], reverse=True)
        topics_text = []
        
        for name, info in sorted_topics:
            files = sorted(info['files'])
            summaries = list(set(info['summaries']))  # Deduplicate
            score = info['max_score']
            
            topics_text.append(
                f"**{name}** (Relevance: {score}/10)\n"
                f"- Source Documents: {', '.join(files)}\n"
                f"- Key Points:\n" + 
                '\n'.join([f'  - {s}' for s in summaries])
            )

        topics_block = '\n\n'.join(topics_text)

        prompt = f"""
Create a consolidated tariff analysis report for Macy's executive team using these extracted topics:

{topics_block}

Guidelines:
1. Start with executive summary highlighting top 3 issues
2. For each topic:
   - Combine insights from different documents
   - Specify source documents
   - Include numerical data/timelines
   - Explain business impact using retail examples
3. Sort by descending relevance score
4. Markdown formatting with clear headings
5. Add "Need Further Review" section for uncaptured documents{content_note}

Avoid technical jargon. Focus on operational impacts and strategic recommendations.
"""
        response = self.client.chat.completions.create(
            model=self.deployment_id,
            messages=[
                {"role": "system", "content": "Senior analyst creating consolidated trade report"},
                {"role": "user", "content": prompt}
            ],
            temperature=0.1
        )
        return response.choices[0].message.content.strip()

# ... [Keep the rest of the original code unchanged for Streamlit components and other agents] ...
    
#     def generate_consolidated_summary_from_topics(self, processed_files):
#         """
#         Generate consolidated summary based on extracted topics (with files and scores).
#         """
#         # Build topic map: name -> files, summaries, max score
#         topic_map = {}
#         for filename, filedata in processed_files.items():
#             for topic in filedata['topics']:
#                 name = topic['name']
#                 if name not in topic_map:
#                     topic_map[name] = {
#                         'files': set(),
#                         'summaries': [],
#                         'max_score': 0
#                     }
#                 topic_map[name]['files'].add(filename)
#                 topic_map[name]['summaries'].append(topic['summary'])
#                 topic_map[name]['max_score'] = max(topic_map[name]['max_score'], topic['score'])

#         # Build combined topic list text for prompt
#         topics_text_lines = []
#         # Sort by descending max_score
#         sorted_topics = sorted(topic_map.items(), key=lambda x: x[1]['max_score'], reverse=True)
#         for name, info in sorted_topics:
#             files_list = ", ".join(sorted(info['files']))
#             combined_summary = " ".join(set(info['summaries']))
#             score = info['max_score']
#             topics_text_lines.append(f"- {name} ({files_list}) [Score: {score}]: {combined_summary}")

#         topics_text = "\n".join(topics_text_lines)

#         prompt = f"""
# You are a senior tariff analyst creating a comprehensive and clear consolidated research summary for Macy's CEO team.

# You have been provided with a list of all unique tariff and trade policy topics extracted from multiple documents, including the source files and their preliminary summaries:

# {topics_text}

# Please create a well-structured, detailed summary report that includes:

# - An introduction about the importance of these topics for Macy's operations.
# - For each topic:
#   - A 2-4 sentence detailed explanation covering impact, key facts, timelines, and relevance.
#   - Any insights or strategic recommendations if appropriate.
# - Highlight any cross-topic interactions or overarching trends that emerged.
# - Present the list sorted from highest to lowest relevance score.
# - Use clear headings, bullet points, and markdown formatting for readability.

# Avoid repeating content directly from original files; synthesize and interpret the information to be valuable for Macy's research team.

# Begin now.
# """

#         response = self.client.chat.completions.create(
#             model=self.deployment_id,
#             messages=[
#                 {"role": "system", "content": "Senior tariff analyst consolidating trade topics into a coherent company briefing."},
#                 {"role": "user", "content": prompt}
#             ],
#             temperature=0.0
#         )
#         return response.choices[0].message.content.strip()



    


class EnhancedChatAgent:
    def __init__(self, client, deployment_id):
        self.client = client
        self.deployment_id = deployment_id
        self.question_types = {
            'tariff_impact': 'Analyzing tariff rate changes and their financial impacts',
            'timeline': 'Identifying implementation timelines and deadlines',
            'geopolitical': 'Evaluating geopolitical factors affecting trade policies',
            'comparison': 'Comparing policies across different regions',
            'compliance': 'Checking regulatory compliance requirements',
            'general': 'General analysis on the Question asked by user',
            'greetings': 'Greetings and general inquiries',
        }

    def _classify_question(self, question):
        prompt = f"""
Classify this question into one of these categories: {list(self.question_types.keys())}. 
Return only the category name.

Question: {question}
"""
        response = self.client.chat.completions.create(
            model=self.deployment_id,
            messages=[
                {"role": "system", "content": "Expert question classifier for Macy's analysis"},
                {"role": "user", "content": prompt}
            ],
            temperature=0
        )
        return response.choices[0].message.content.strip().lower()
    

    def _stream_greeting_response(self, question):
        prompt = f"""
You are an senoir assistant responding ONLY to greetings from Macy's executive staff. Reply with a polite greeting message only. No additional information but respond properly for every question do not use any other source.

Greeting input:
{question}

Your reply:
"""
        messages = [
            {"role": "system", "content": "Assistant specialized in greeting responses."},
            {"role": "user", "content": prompt}
        ]
        response = self.client.chat.completions.create(
            model=self.deployment_id,
            messages=messages,
            temperature=0,
            stream=True
        )
        return response
    

    def stream_answer(self, question, files_context):
        q_type = self._classify_question(question)

        if q_type == 'greetings':
            response = self._stream_greeting_response(question)
            return response, q_type, 'Greeting response'
        
        analysis_approach = self.question_types.get(q_type, 'General analysis')

        base_prompt = f"""
You are a senior analyst answering questions for Macy's executive team. Follow these steps:

1. {analysis_approach}
2. Cross-reference all relevant documents
3. Identify numerical data and timelines
4. Assess impacts on retail operations
5. Formulate executive-level recommendations
6. If User ask from a single document, selecting multiple documents give them answer from only the document which is asked for.
7. If User ask from multiple documents, selecting multiple documents give them answer from all the documents.
8. If User ask a question like greetings, answer it with "Hello, how can I help you?" Nothing else from the documents.


Current question type: {q_type.upper()}

Question: {question}

Relevant documents:
{files_context}

Provide a structured response with:
- Clear headings for each section
- Bullet points for key findings
- Bolded key figures and dates
- Separate recommendations section

If information conflicts between documents, note this explicitly.
"""
        
        messages = [
            {"role": "system", "content": "Senior analyst synthesizing information from multiple documents"},
            {"role": "user", "content": base_prompt}
        ]

        response = self.client.chat.completions.create(
            model=self.deployment_id,
            messages=messages,
            temperature=0,
            stream=True
        )
        return response, q_type, analysis_approach


# ===== Streamlit App Main =====

def main():
    st.set_page_config(page_title="Tariff & Trade Analyst Pro", layout="wide")
    st.title("📄 Macy's Tariff & Trade Policy Analyzer")

    # Initialize session state vars
    if "session" not in st.session_state:
        st.session_state.session = {
            'uploaded_files': {},
            'processed_files': {},
            'active_file': None,
            'active_topic': None,
            'topic_details_cache': {},
            'chat_history': [],
            'consolidated_summary': None
        }
    
    if "client" not in st.session_state:
        st.session_state.client = openai.AzureOpenAI(
            api_key=os.getenv("AZURE_API_KEY"),
            azure_endpoint=os.getenv("AZURE_ENDPOINT"),
            api_version="2023-12-01-preview"
        )
    
    if "agent" not in st.session_state:
        deployment = os.getenv("AZURE_DEPLOYMENT")
        st.session_state.agent = TariffTradeAgent(st.session_state.client, deployment)
        st.session_state.chat_agent = EnhancedChatAgent(st.session_state.client, deployment)

    # File upload section
    st.subheader("1. Upload Documents")
    uploaded_files = st.file_uploader(
        "Upload Macy's earnings transcript PDFs",
        type=["pdf"],
        accept_multiple_files=True,
        key="file_uploader"
    )

    # Store uploaded files in session state
    if uploaded_files:
        for file in uploaded_files:
            if file.name not in st.session_state.session['uploaded_files']:
                st.session_state.session['uploaded_files'][file.name] = file.getvalue()

    # Processing button
    st.subheader("2. Analyze Documents")
    col1, col2 = st.columns([1, 3])
    with col1:
        process_btn = st.button(
            "🔍 Generate Tariff Topics & Summary",
            help="Analyze uploaded documents for tariff-related topics and generate consolidated summary",
            type="primary"
        )
    
    if process_btn:
        if not st.session_state.session['uploaded_files']:
            st.error("Please upload documents first")
            st.stop()

        with st.spinner("Analyzing documents for tariff topics..."):
            for filename, file_bytes in st.session_state.session['uploaded_files'].items():
                if filename not in st.session_state.session['processed_files']:
                    try:
                        # Convert bytes to file-like object
                        file_obj = BytesIO(file_bytes)
                        file_obj.name = filename
                        
                        text = extract_text(file_obj)
                        if text:
                            topics = st.session_state.agent.analyze_document(text)
                            st.session_state.session['processed_files'][filename] = {
                                'text': text,
                                'topics': topics,
                                'context': text
                            }
                            st.toast(f"✅ Processed {filename}")
                        else:
                            st.error(f"❌ Failed to process {filename}")
                    except Exception as e:
                        st.error(f"Error processing {filename}: {str(e)}")
            st.success("Analysis completed! Generating consolidated summary...")

        # Generate consolidated summary **immediately** after processing
        with st.spinner("Generating consolidated summary..."):
            summary_text = st.session_state.agent.generate_consolidated_summary_from_topics(
                st.session_state.session['processed_files']
            )
            st.session_state.session['consolidated_summary'] = summary_text

        st.success("Consolidated summary generated!")

    # Show consolidated summary if available
    if st.session_state.session.get('consolidated_summary'):
        st.header("🌟 Consolidated Tariff & Trade Policy Summary")
        st.markdown(st.session_state.session['consolidated_summary'])
        st.markdown("---")

    # Documents & Topics navigation
    st.sidebar.header("Documents & Topics")
    for filename, filedata in st.session_state.session['processed_files'].items():
        expanded = filename == st.session_state.session['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:
                is_selected = (st.session_state.session['active_file'] == filename and
                               st.session_state.session['active_topic'] == topic['name'])
                style = "primary" if is_selected else "secondary"
                if st.button(f"• {topic['name']} (Score: {topic['score']})",
                             key=f"sidebar_topic_{filename}_{topic['name']}",
                             type=style,
                             use_container_width=True):
                    if (st.session_state.session['active_file'] != filename or
                        st.session_state.session['active_topic'] != topic['name']):
                        st.session_state.session['active_file'] = filename
                        st.session_state.session['active_topic'] = topic['name']

    # Detailed topic analysis
    st.subheader("Detailed Topic Analysis")
    sess = st.session_state.session
    if sess['active_file'] and sess['active_topic']:
        key_cache = (sess['active_file'], sess['active_topic'])
        if key_cache in sess['topic_details_cache']:
            detail = sess['topic_details_cache'][key_cache]
            st.markdown(detail)
        else:
            text = sess['processed_files'][sess['active_file']]['context']
            topic = sess['active_topic']

            placeholder = st.empty()
            collected_text = ""

            try:
                response = st.session_state.agent.stream_generate_detail(text, topic)
                for chunk in response:
                    if hasattr(chunk, "choices") and chunk.choices:
                        delta = chunk.choices[0].delta
                        content = getattr(delta, "content", "")
                        if content:
                            collected_text += content
                            placeholder.markdown(collected_text + "▌")
                placeholder.markdown(collected_text)
                sess['topic_details_cache'][key_cache] = collected_text
            except Exception as e:
                st.error(f"Error generating detail: {e}")
    else:
        st.info("Select a file and topic from the sidebar to view detailed analysis.")

    # Enhanced Chat Section
    st.divider()
    st.subheader("💬 Advanced Document Analysis")

    selected_files = st.multiselect(
        "Select documents for analysis:",
        options=list(st.session_state.session['processed_files'].keys()),
        default=[]
    )

    # Display chat history
    for msg in sess['chat_history']:
        with st.chat_message(msg["role"]):
            if msg["role"] == "assistant":
                if 'analysis_type' in msg:
                    st.markdown(f"**Analysis Type:** {msg['analysis_type']}")
                    st.markdown(f"**Process:** {msg['process']}")
                    st.markdown("**Findings:**")
                    st.markdown(msg["content"])
                else:
                    st.markdown(msg["content"])
            else:
                st.markdown(msg["content"])

    if prompt := st.chat_input("Ask a complex question about tariffs and trade policies"):
        if not selected_files:
            st.error("Please select one or more documents before asking a question.")
            st.stop()

        # Add user question to history
        sess['chat_history'].append({
            "role": "user",
            "content": prompt,
            "timestamp": datetime.utcnow().isoformat()
        })

        # Prepare combined context
        files_context = "\n\n".join([
            f"=== {filename} ===\n{sess['processed_files'][filename]['context']}"
            for filename in selected_files
        ])

        # Create containers for the response
        response_container = st.chat_message("assistant")
        status_container = response_container.status("Beginning analysis...")
        response_placeholder = response_container.empty()

        full_response = ""
        
        try:
            # Get streaming response
            response_stream, q_type, analysis_approach = st.session_state.chat_agent.stream_answer(
                prompt, files_context
            )

            # Update status with analysis type
            status_container.update(
                label=f"Analysis Type: {q_type.upper()}",
                state="complete",
                expanded=True
            )

            # Display process information
            with response_container.expander("Analysis Process Details"):
                st.markdown(f"""
                **Documents Analyzed:**
                {chr(10).join(['- '+f for f in selected_files])}
                """)

            # Stream the response
            for chunk in response_stream:
                if hasattr(chunk, "choices") and chunk.choices:
                    delta = chunk.choices[0].delta
                    content = getattr(delta, "content", "")
                    if content:
                        full_response += content
                        response_placeholder.markdown(full_response + "▌")

            # Finalize response formatting
            response_placeholder.markdown(full_response)

            # Add to chat history with metadata
            sess['chat_history'].append({
                "role": "assistant",
                "content": full_response,
                # "analysis_type": q_type.upper(),
                # "process": analysis_approach,
                "timestamp": datetime.utcnow().isoformat()
            })

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

    # Reset session button
    if st.button("🔄 Reset Session", type="secondary"):
        preserved_keys = ["client", "agent", "chat_agent"]
        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': {},
            'active_file': None,
            'active_topic': None,
            'topic_details_cache': {},
            'chat_history': [],
            'consolidated_summary': None
        }
        st.experimental_rerun()


if __name__ == "__main__":
    main()


i need similar topics like Tarrifs 
# ===== SustainabilityAgent =====
class SustainabilityAgent:
    def __init__(self, client, deployment_id):
        self.client = client
        self.deployment_id = deployment_id
        
    def analyze_document(self, text):
        prompt = """
You are a senior analyst specializing in Sustainability and ESG (Environmental, Social, and Governance) topics. Identify all significant ESG-related topics mentioned in the text that could impact Macy's corporate responsibility, supply chain ethics, or brand reputation.

For each topic, provide:
- A clear topic name.
- A concise 1-sentence summary describing its impact or relevance.
- A relevance score from 1 (minor mention) to 10 (critical detail).

Output one topic per line with format:

- Topic Name: Summary [Relevance Score]

Examples:
- Carbon Emissions Reduction: Macy's aims to cut carbon footprint 30% by 2030 [9]
- Ethical Sourcing Practices: Strengthening supplier labor standards [8]
- Renewable Energy Use: Increasing use of solar energy in stores [7]

Ignore content unrelated to ESG or sustainability.

Begin analysis. Text:
"""
        response = self.client.chat.completions.create(
            model=self.deployment_id,
            messages=[
                {"role": "system", "content": "You analyze sustainability and ESG issues for Macy's management team."},
                {"role": "user", "content": prompt + "\n\n" + text}
            ],
            temperature=0
        )
        return self._parse_topics(response.choices[0].message.content.strip())
    
    def _parse_topics(self, text):
        topics = []
        pattern = r'^- (.*?): (.*?) \[(\d+)\]$'
        for line in text.split('\n'):
            line = line.strip()
            match = re.match(pattern, line)
            if match:
                topics.append({
                    'name': match.group(1).strip(),
                    'summary': match.group(2).strip(),
                    'score': int(match.group(3).strip())
                })
        return topics

    def stream_generate_detail(self, text, topic):
        prompt = f"""
You are drafting a detailed briefing for Macy's CEO on the sustainability topic '{topic}' extracted from the document.

Provide:
1. An overview.
2. Key insights and direct excerpts.
3. Relevant figures or timelines.
4. Potential impact on Macy's CSR, brand, and supply chain.
5. 3-5 strategic recommendations.

Document:
{text}

Topic: {topic}

Deliver as structured markdown for executives.
"""
        messages = [
            {"role": "system", "content": "Senior ESG analyst briefing to CEO"},
            {"role": "user", "content": prompt}
        ]
        response = self.client.chat.completions.create(
            model=self.deployment_id,
            messages=messages,
            temperature=0,
            stream=True,
        )
        return response

    def generate_consolidated_summary_from_topics(self, processed_files):
        topic_map = {}
        for filename, filedata in processed_files.items():
            for topic in filedata['topics']:
                name = topic['name']
                if name not in topic_map:
                    topic_map[name] = {
                        'files': set(),
                        'summaries': [],
                        'max_score': 0,
                    }
                topic_map[name]['files'].add(filename)
                topic_map[name]['summaries'].append(topic['summary'])
                topic_map[name]['max_score'] = max(topic_map[name]['max_score'], topic['score'])

        topics_text_lines = []
        sorted_topics = sorted(topic_map.items(), key=lambda x: x[1]['max_score'], reverse=True)
        for name, info in sorted_topics:
            files_list = ", ".join(sorted(info['files']))
            combined_summary = " | ".join(set(info['summaries']))
            score = info['max_score']
            topics_text_lines.append(f"- {name} (Files: {files_list}) [Score: {score}]: {combined_summary}")

        topics_text = "\n".join(topics_text_lines)
        prompt = f"""
You are a senior sustainability analyst preparing an executive summary report for Macy's CEO team.

Topics extracted from documents:
{topics_text}

Please provide:
- An introduction on the importance of these sustainability and ESG topics.
- For each topic: a 2-4 sentence detailed explanation with impacts, key insights, and strategic recommendations.
- Highlight any overarching trends across topics.
- Use clear markdown formatting with headings and bullet points.
- Sort topics from highest to lowest by relevance.

Begin:
"""
        response = self.client.chat.completions.create(
            model=self.deployment_id,
            messages=[
                {"role": "system", "content": "Senior ESG analyst summarizing sustainability topics for Macy's management."},
                {"role": "user", "content": prompt},
            ],
            temperature=0,
        )
        return response.choices[0].message.content.strip()

mak ea button for this and i need to display directly as for tarrifs we are showing we can directly show the tarrif topics in sidebar but when clicked on Sustainability  we can directly see the only Sustainability  topics and the consolidate genrated summary and all

give me full code do not left any other funtinality of my code