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 ""

import re

# ===== TariffTradeAgent =====
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.  
Carefully identify all significant tariff, trade, and related government policy topics in the text that could affect Macy's business operations, including supply chain, procurement, retail sales, and technology.

For each topic, provide:  
- A clear and concise topic name (e.g., "Vietnam Production Impact", "US-China Tariffs").  
- A brief 1-sentence summary describing its operational or strategic impact on Macy's retail business.  
- A relevance score from 1 (minor) to 10 (critical).  
- Specific examples of how this issue may influence US retail, consumer behavior, procurement, supply chain logistics, or merchandising strategy.

Format your response as one topic per line in this format:  
- Topic Name: Summary [Relevance Score]

Example lines:  
- Vietnam Factory Shutdown: Supply chain delays expected due to production halt in Vietnam [9]  
- Tariff Changes on Apparel: Increased costs on imports from China may raise prices [8]

Analyze the following text for all relevant tariff and trade topics. Text:
"""
        response = self.client.chat.completions.create(
            model=self.deployment_id,
            messages=[
                {"role": "system", "content": "Tariff and trade policy 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 stream_generate_detail(self, text, topic):
        prompt = f"""You are preparing a detailed briefing for Macy's CEO on the tariff and trade policy topic '{topic}', based on the provided document.

Provide:
1. A clear overview of the topic.
2. Key insights and any relevant direct excerpts from the document.
3. Important figures, dates, or timelines if available.
4. Detailed analysis on how this topic could impact Macy's retail operations, supply chain, procurement, and consumer behavior.
5. 3-5 actionable strategic recommendations for Macy's leadership.

Document:
{text}

Topic: {topic}

Please deliver the response as structured markdown suitable for executive briefing."""
        messages = [
            {"role": "system", "content": "Senior tariff analyst briefing to Macy's 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 = {}
        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)

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

        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']))
            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 and trade policy analysis report for Macy's executive team based on these topics:

{topics_block}

Instructions:
1. Begin with an executive summary highlighting the top 3 most impactful issues.
2. For each topic:
   - Synthesize insights from all source documents.
   - Specify supporting document names.
   - Include any numerical data or timelines.
   - Explain the business impact with examples tailored to Macy's retail operations.
   - Detailed analysis on how this topic could impact Macy's retail operations, supply chain, procurement, and consumer behavior.
   - 3-5 actionable strategic recommendations for Macy's leadership.
3. Order topics by descending relevance score.
4. Use markdown formatting with clear headings.
5. Add a "Need Further Review" section for documents without extracted topics.{content_note}

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


# ===== 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, ESG (Environmental, Social, Governance), and corporate responsibility topics.  
Identify all significant ESG and sustainability-related topics in the text that could influence Macy's corporate responsibility, brand reputation, ethical supply chain management, or stakeholder relations.

For each topic, provide:  
- A concise and clear topic name (e.g., "Carbon Footprint Reduction", "Ethical Sourcing Initiatives").  
- A brief 1-sentence summary describing the topic's relevance or impact on Macy's operations or reputation.  
- A relevance score from 1 (minor) to 10 (critical).

Format the output as one topic per line in this format:  
- Topic Name: Summary [Relevance Score]

Example lines:  
- Renewable Energy Adoption: Macy's increasing solar panel use in stores nationwide [7]  
- Labor Rights Compliance: Strengthening oversight on supplier labor conditions [9]

Analyze the following text for relevant ESG and sustainability topics. Text:
"""
        response = self.client.chat.completions.create(
            model=self.deployment_id,
            messages=[
                {"role": "system", "content": "Sustainability and ESG analyst for Macy's management"},
                {"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}', based on the provided document.

Provide:
1. An overview of the topic.
2. Key insights and relevant excerpts.
3. Current figures, goals, or timelines.
4. Potential impact on Macy's brand, CSR, supply chain ethics, and stakeholder relations.
5. 3-5 strategic recommendations to enhance Macy's sustainability performance.

Document:
{text}

Topic: {topic}

Deliver as structured markdown for executive review."""
        messages = [
            {"role": "system", "content": "Senior ESG analyst preparing briefing for Macy's 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 = {}
        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)

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

        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']))
            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 sustainability and ESG report for Macy's executive team based on these topics:

{topics_block}

Guidelines:
1. Start with an executive summary highlighting top 3 ESG priorities.
2. For each topic:
   - Combine insights from all documents.
   - Include source document references.
   - Provide relevant numerical targets or timelines.
   - Explain impact on Macy's reputation and operational ethics with specific retail examples.
3. Sort topics by descending relevance score.
4. Format using markdown with clear headings.
5. Add a section "Need Further Review" for documents without extracted topics.{content_note}

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


# ===== EnhancedChatAgent (unchanged) =====
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 Document Analysis & Chat Tool")

    # Initialize session state vars
    if "session" not in st.session_state:
        st.session_state.session = {
            'uploaded_files': {},
            'processed_files_tariff': {},
            'processed_files_sustainability': {},
            'active_file': None,
            'active_topic': None,
            'topic_details_cache': {},  # (mode, file, topic)
            'chat_history': [],
            'consolidated_summary_tariff': None,
            'consolidated_summary_sustainability': None,
            'last_analysis_mode': 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_tariff" not in st.session_state:
        deployment = os.getenv("AZURE_DEPLOYMENT")
        st.session_state.agent_tariff = TariffTradeAgent(st.session_state.client, deployment)
        st.session_state.agent_sustainability = SustainabilityAgent(st.session_state.client, deployment)
        st.session_state.chat_agent = EnhancedChatAgent(st.session_state.client, deployment)

    sess = st.session_state.session

    # File Upload Section (shared)
    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"
    )
    if uploaded_files:
        for file in uploaded_files:
            if file.name not in sess['uploaded_files']:
                sess['uploaded_files'][file.name] = file.getvalue()

    # Analyze Documents Buttons
    st.subheader("2. Analyze Documents")
    col1, col2 = st.columns(2)

    analyze_tariff_clicked = col1.button("🔍 Analyze Tariffs & Trade", type="primary")
    analyze_sustain_clicked = col2.button("♻️ Analyze Sustainability & ESG", type="primary")

    if analyze_tariff_clicked or analyze_sustain_clicked:
        if not sess['uploaded_files']:
            st.error("Please upload documents before analyzing.")
            st.stop()

        if analyze_tariff_clicked:
            mode = 'tariff'
            agent = st.session_state.agent_tariff
            processed_files = sess['processed_files_tariff']
        else:
            mode = 'sustainability'
            agent = st.session_state.agent_sustainability
            processed_files = sess['processed_files_sustainability']

        with st.spinner(f"Analyzing documents for {mode} topics..."):
            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)
                            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)}")

        with st.spinner("Generating consolidated summary..."):
            summary_text = agent.generate_consolidated_summary_from_topics(processed_files)
            if mode == 'tariff':
                sess['consolidated_summary_tariff'] = summary_text
            else:
                sess['consolidated_summary_sustainability'] = summary_text

        sess['last_analysis_mode'] = mode
        sess['active_file'] = None  # reset active topic selections when mode changes
        sess['active_topic'] = None

        st.success(f"Analysis & summary generated for {mode}!")

    # Determine which mode to display (last button clicked)
    mode_to_display = sess['last_analysis_mode']
    if mode_to_display is None:
        st.info("Click a button above to analyze Tariffs/Trade or Sustainability topics.")
        return

    # Pick relevant processed_files dict & consolidated summary
    if mode_to_display == 'tariff':
        processed_files = sess['processed_files_tariff']
        consolidated_summary = sess['consolidated_summary_tariff']
        sidebar_header = "Documents & Tariff Topics"
        main_header = "🌟 Consolidated Tariff & Trade Summary"
        agent = st.session_state.agent_tariff
    else:
        processed_files = sess['processed_files_sustainability']
        consolidated_summary = sess['consolidated_summary_sustainability']
        sidebar_header = "Documents & Sustainability Topics"
        main_header = "🌟 Consolidated Sustainability & ESG Summary"
        agent = st.session_state.agent_sustainability

    # Show consolidated summary
    st.header(main_header)
    if consolidated_summary:
        st.markdown(consolidated_summary)
    else:
        st.info(f"No consolidated summary available yet for {mode_to_display.capitalize()}. Please analyze documents.")

    st.markdown("---")

    # Sidebar: Documents and topics
    st.sidebar.header(sidebar_header)
    for filename, filedata in processed_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:
                is_selected = (sess['active_file'] == filename and sess['active_topic'] == topic['name'])
                style = "primary" if is_selected else "secondary"
                if st.button(  # CHANGED FROM st.sidebar.button
                    f"• {topic['name']} (Score: {topic['score']})",
                    key=f"{mode_to_display}_topic_{filename}_{topic['name']}",
                    type=style,
                    use_container_width=True
                ):
                    sess['active_file'] = filename
                    sess['active_topic'] = topic['name']

    # Detailed Topic Analysis
    st.subheader("Detailed Topic Analysis")
    if sess['active_file'] and sess['active_topic']:
        key_cache = (mode_to_display, 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 = processed_files[sess['active_file']]['context']
            topic = sess['active_topic']

            placeholder = st.empty()
            collected_text = ""

            try:
                response = 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.")

    st.divider()

    # Advanced Chat Section (Tariff only)
    st.subheader("💬 Advanced Document Analysis (Tariff & Trade)")

    tariff_files = list(sess['processed_files_tariff'].keys())
    selected_files = st.multiselect(
        "Select documents for analysis:",
        options=tariff_files,
        default=[]
    )

    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()

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

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

        response_container = st.chat_message("assistant")
        status_container = response_container.status("Beginning analysis...")
        response_placeholder = response_container.empty()

        full_response = ""

        try:
            response_stream, q_type, analysis_approach = st.session_state.chat_agent.stream_answer(
                prompt, files_context
            )
            status_container.update(label=f"Analysis Type: {q_type.upper()}", state="complete", expanded=True)

            with response_container.expander("Analysis Process Details"):
                st.markdown(f"**Documents Analyzed:**\n" + '\n'.join(['- '+f for f in selected_files]))

            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 + "▌")

            response_placeholder.markdown(full_response)

            sess['chat_history'].append({
                "role": "assistant",
                "content": full_response,
                "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_tariff", "agent_sustainability", "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_tariff': {},
            'processed_files_sustainability': {},
            'active_file': None,
            'active_topic': None,
            'topic_details_cache': {},
            'chat_history': [],
            'consolidated_summary_tariff': None,
            'consolidated_summary_sustainability': None,
            'last_analysis_mode': None
        }
        st.experimental_rerun()


if __name__ == "__main__":
    main()


for consolidate summary you can doesnot send the topics but you csn send the files content and get the relevent summay from the files also make it accrodingly as we are finding our topics

give me full code do not change any thing else