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.
- A concise 1-sentence summary describing its impact or relevance.
- A relevance score: 1 (minor mention) to 10 (critical detail).

If possible, indicate whether the impact is generally an opportunity, risk, or neutral factor, but do NOT explicitly categorize with labels like [Positive] or [Negative].

Output one topic per line, format:

- Topic Name: Summary [Relevance Score]

Example:
- USMCA Trade Agreement: Expected tariff reductions improve apparel import costs for Macy's [8]
- Chinese Tariffs Impact: Increased tariffs lead to higher electronic goods prices [7]
- Implementation Timeline: Changes will roll out over a 2-year phased schedule [6]

Ignore content unrelated to tariffs or international trade policies.

Begin analysis. Text:
"""
        response = self.client.chat.completions.create(
            model=self.deployment_id,
            messages=[
                {"role": "system", "content": "You analyze tariffs & trade policy for Macy's CEO research 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 business briefing for the CEO of Macy's on the tariff/trade topic '{topic}' extracted from the following document. 

Provide:
1. An overview of the topic.
2. Key insights with direct excerpts.
3. Relevant numerical or timeline data.
4. Potential impacts on Macy's supply chain & sales.
5. Strategic recommendations in 3-5 concise points.

Document:
{text}

Topic: {topic}

Deliver the response as a clear and structured markdown for easy reading by senior management.
"""
        messages = [
            {"role": "system", "content": "Senior tariff 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):
        """
        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()

    def generate_consolidated_summary(self, combined_file_texts):
        """Legacy method kept for compatibility - uses raw texts, not topics."""
        # (You can remove or keep this if you want)
        prompt = """
You are a senior tariff analyst creating a single consolidated summary of all tariff and trade policy topics extracted from multiple documents. Some topics may appear in multiple documents. For each unique topic, provide:

- The topic name.
- A detailed 2-4 sentence paragraph describing its impact, key facts, timelines, and relevance to Macy's operations and strategy.
- A relevance score from 1 (minor mention) to 10 (critical detail).
- The file(s) from which this topic was drawn.

Present the list sorted by descending relevance score, format each entry as:

- Topic Name (Source file(s)) [Score]:
Detailed paragraph of the topic here.

Example:

- USMCA Trade Agreement (File1.pdf, File3.pdf) [8]:
The USMCA agreement significantly reduces tariffs on apparel imports that Macy's relies on, potentially lowering costs and improving supply chain efficiency. Changes will be phased in over 3 years, with opportunities for renegotiation.

- Chinese Tariffs Impact (File2.pdf) [7]:
Tariffs imposed on Chinese imports have increased costs of electronics sold by Macy's. The increase is expected to remain for at least 2 years, forcing pricing adjustments.

Do not include unrelated content.

Begin now. Texts by file below:
""" + combined_file_texts

        response = self.client.chat.completions.create(
            model=self.deployment_id,
            messages=[
                {"role": "system", "content": "Senior tariff analyst consolidating trade policy topics for CEO summary."},
                {"role": "user", "content": prompt}
            ],
            temperature=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()

generate_consolidated_summary for this we have to use the first agent outputs and with that we can send all the files too to get the output only for that as for Summarizing them making it more better results