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(self, combined_file_texts):
        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_answer(self, question, files_context):
        q_type = self._classify_question(question)
        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",
            help="Analyze uploaded documents for tariff-related topics",
            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! Select documents in the sidebar to view topics.")

    # Consolidated summary section
    st.sidebar.markdown("## Consolidated Summary")
    if st.sidebar.button("Generate Consolidated Summary"):
        if not st.session_state.session['processed_files']:
            st.sidebar.error("No processed documents available")
            st.stop()

        combined_texts = ""
        for filename, filedata in st.session_state.session['processed_files'].items():
            combined_texts += f"\n---\nFile: {filename}\n\n{filedata['context']}\n\n"
        
        with st.spinner("Generating consolidated summary across all files..."):
            summary_text = st.session_state.agent.generate_consolidated_summary(combined_texts)
            st.session_state.session['consolidated_summary'] = summary_text

    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"""
                **Question Type:** {q_type}
                **Analysis Approach:** {analysis_approach}
                **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.rerun() 


if __name__ == "__main__":
    main()

remove the consise genrated button and genrate the consise summary after the all topics genrated from the files make sure to include only those topics which first agent had maked