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

# ===== 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.  
You have received a transcript/document from another company.  

Carefully identify all significant tariff, trade, and related government policy topics and detailed briefing 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 describing in 50-75 words including its operational or strategic impact on Macy's retail business, considering the external company's context also include A clear overview of the topic, Key insights and any relevant direct excerpts from the document, Important figures, dates, or timelines if available, Detailed analysis on how this topic could impact Macy's retail operations, supply chain, procurement, and consumer behavior.  
- A relevance score from 1 (minor) to 10 (critical).  

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

Example:  
- 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 external company document 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
        )
        raw_text = response.choices[0].message.content.strip()
        return self._parse_topics(raw_text)

    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 external company 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):

        topics_block = ""
        for filename, filedata in processed_files.items():
            topics_block += f"### Topics from document: {filename}\n"
            for topic in filedata['topics']:
                topics_block += f"- **{topic['name']}**: {topic['summary']} (Relevance: {topic['score']})\n"
            topics_block += "\n"

        prompt = f"""You are a senior tariff and trade analyst preparing a consolidated tariff and trade policy report for Macy's executive team.

You have reviewed tariff and trade topics extracted from multiple external company documents. These topics include their names, summaries, and relevance scores. Based on their content, please:

1. Identify and describe from all documents all their impactful tariff, trade, and related government policy topics that could affect Macy's business operations, including supply chain, procurement, retail sales, and technology.
2. For each topic:
   - Provide a concise summary from combining all the documents.
   - Include relevant examples, numerical data, or timelines from the documents also provide from which document it is taken.
   - Explain operational and strategic impacts—positive, negative, or neutral—specifically tailored to Macy's retail business.
   - Provide 3-5 actionable strategic recommendations for Macy's leadership.
3. Use markdown formatting with clear headings.
4. Avoid jargon and focus on practical operational impacts and strategy insights.

Format your response in markdown with clear headings for each topic.

Here are the topics extracted from documents:

{topics_block}

Deliver a detailed, synthesized, and actionable consolidated tariff and trade policy report for Macy's.
"""

        response = self.client.chat.completions.create(
            model=self.deployment_id,
            messages=[
                {"role": "system", "content": "Senior tariff and trade analyst consolidating extracted topics for Macy's"},
                {"role": "user", "content": prompt}
            ],
            temperature=0,
        )
        return response.choices[0].message.content.strip()

# ===== SalesAgent =====
class SalesAgent:
    def __init__(self, client, deployment_id):
        self.client = client
        self.deployment_id = deployment_id

    def analyze_document(self, text):
        prompt = """You are a senior sales analyst specializing in retail sales and consumer behavior, assisting Macy's executive team.  
You have received a transcript/document from another company.

Identify all significant sales and consumer behavior topics that could affect Macy's retail sales strategy, revenue, marketing effectiveness, and customer engagement.

For each topic, provide:  
- A concise topic name (e.g., "Holiday Sales Surge", "Online Shopping Trends").  
- A brief describing in 50-75 words including its relevance or impact on Macy's sales, considering the external company's context, also include An overview of the topic, Key insights and relevant excerpts, Important figures, dates, or timelines, Detailed analysis of the impact on Macy's retail sales, marketing, and customer engagement. 
- A relevance score from 1 (minor) to 10 (critical).

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

Example:  
- Holiday Sales Boost: Significant increase in sales volume expected during holidays [9]  
- E-commerce Growth: Online shopping rise impacting in-store sales [8]

Analyze the following external company document:
"""
        response = self.client.chat.completions.create(
            model=self.deployment_id,
            messages=[
                {"role": "system", "content": "Sales analyst for Macy's executive team"},
                {"role": "user", "content": prompt + "\n\n" + text}
            ],
            temperature=0,
        )
        raw_text = response.choices[0].message.content.strip()
        return self._parse_topics(raw_text)

    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 sales topic '{topic}', based on the external company document.

Provide:
1. An overview of the topic.
2. Key insights and relevant excerpts.
3. Important figures, dates, or timelines.
4. Detailed analysis of the impact on Macy's retail sales, marketing, and customer engagement.
5. 3-5 actionable recommendations for Macy's leadership.

Document:
{text}

Topic: {topic}

Please provide structured markdown briefing.
"""
        messages = [
            {"role": "system", "content": "Senior sales analyst 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):
        topics_block = ""
        for filename, filedata in processed_files.items():
            topics_block += f"### Topics from document: {filename}\n"
            for topic in filedata['topics']:
                topics_block += f"- **{topic['name']}**: {topic['summary']} (Relevance: {topic['score']})\n"
            topics_block += "\n"

        prompt = f"""You are a senior sales analyst preparing a consolidated sales and consumer behavior report for Macy's executive team.

You have reviewed sales and consumer behavior topics extracted from multiple external company documents. These topics include their names, summaries, and relevance scores. Based on their content, please:


1. Identify and describe from all documents all their impactful sales and consumer behavior topics that could affect Macy's retail sales, marketing, and customer engagement.
2. For each topic:
   - Provide a concise summary from combining all the documents.
   - Include relevant numerical data or timelines from the documents also provide from which document it is taken.
   - Explain operational and strategic impacts—positive, negative, or neutral—specifically tailored to Macy's retail business.
   - Provide 3-5 actionable strategic recommendations for Macy's leadership.
3. Use markdown formatting with clear headings.
4. Avoid jargon and focus on actionable insights tailored for Macy's.

Format response in markdown with clear headings.

Extracted topics from documents:

{topics_block}

Deliver a detailed, actionable consolidated sales report for Macy's.
"""

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


# ===== SupplyChainAgent =====
class SupplyChainAgent:
    def __init__(self, client, deployment_id):
        self.client = client
        self.deployment_id = deployment_id

    def analyze_document(self, text):
        prompt = """You are a senior supply chain analyst specializing in logistics and procurement for Macy's operations.  
You have received a transcript/document from another company.

Identify all significant supply chain and logistics topics in the text that could affect Macy's procurement, inventory, delivery timelines, and supplier relationships.

For each topic, provide:  
- A concise topic name (e.g., "Port Congestion", "Supplier Delays").  
- A brief describing in 50-75 words including its impact on Macy's supply chain operations also include Overview of the topic, Key insights and relevant excerpts, Important figures, timelines, or operational metrics, Detailed analysis on impact to Macy's supply chain, procurement, and logistics.  
- A relevance score from 1 (minor) to 10 (critical).

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

Example lines:  
- Port Congestion: Delays expected affecting inventory availability [9]  
- Supplier Delay: Lead times extended due to supplier issues [8]

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

    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 supply chain topic '{topic}', based on the external company document.

Provide:
1. Overview of the topic.
2. Key insights and relevant excerpts.
3. Important figures, timelines, or operational metrics.
4. Detailed analysis on impact to Macy's supply chain, procurement, and logistics.
5. 3-5 actionable recommendations for Macy's leadership.

Document:
{text}

Topic: {topic}

Please provide structured markdown briefing.
"""
        messages = [
            {"role": "system", "content": "Senior supply chain analyst 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):
        topics_block = ""
        for filename, filedata in processed_files.items():
            topics_block += f"### Topics from document: {filename}\n"
            for topic in filedata['topics']:
                topics_block += f"- **{topic['name']}**: {topic['summary']} (Relevance: {topic['score']})\n"
            topics_block += "\n"

        prompt = f"""You are a senior supply chain analyst preparing a consolidated supply chain and procurement report for Macy's executive team.

You have reviewed supply chain and logistics topics extracted from multiple external company documents. These topics include their names, summaries, and relevance scores. Based on their content, please:

1. Identify and describe all their impactful supply chain and logistics topics that could affect Macy's inventory management, procurement, delivery timelines, and supplier relationships.
2. For each topic:
   - Provide a concise summary from combining all the documents.
   - Include relevant metrics, timelines, or examples from the documents also provide from which document it is taken.
   - Explain operational and strategic impacts—positive, negative, or neutral—specifically tailored to Macy's supply chain and logistics.
   - Provide 3-5 actionable strategic recommendations for Macy's leadership.
3. Use markdown formatting with clear headings.
4. Avoid jargon and focus on practical operational impacts and strategy insights.

Format response in markdown with clear headings.

Extracted topics from documents:

{topics_block}

Deliver a detailed, actionable consolidated supply chain report for Macy's.
"""

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

# ===== EnhancedChatAgent =====
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 vars
    if "session" not in st.session_state:
        st.session_state.session = {
            'uploaded_files': {},
            'processed_files_tariff': {},
            'processed_files_sales': {},
            'processed_files_supplychain': {},
            'active_file': None,
            'active_topic': None,
            'topic_details_cache': {},
            'chat_history': [],
            'consolidated_summary_tariff': None,
            'consolidated_summary_sales': None,
            'consolidated_summary_supplychain': 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_sales = SalesAgent(st.session_state.client, deployment)
        st.session_state.agent_supplychain = SupplyChainAgent(st.session_state.client, deployment)
        st.session_state.chat_agent = EnhancedChatAgent(st.session_state.client, deployment)

    sess = st.session_state.session

    # File upload
    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 buttons
    st.subheader("2. Analyze Documents")
    col1, col2, col3 = st.columns(3)

    analyze_tariff_clicked = col1.button("🔍 Analyze Tariffs & Trade", type="primary")
    analyze_sales_clicked = col2.button("📈 Analyze Sales & Consumer Behavior", type="primary")
    analyze_supplychain_clicked = col3.button("🚚 Analyze Supply Chain & Logistics", type="primary")

    mode = None
    agent = None
    processed_files = None

    if analyze_tariff_clicked  or analyze_sales_clicked or analyze_supplychain_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']
        elif analyze_sales_clicked:
            mode = 'sales'
            agent = st.session_state.agent_sales
            processed_files = sess['processed_files_sales']
        elif analyze_supplychain_clicked:
            mode = 'supplychain'
            agent = st.session_state.agent_supplychain
            processed_files = sess['processed_files_supplychain']

        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
            elif mode == 'sales':
                sess['consolidated_summary_sales'] = summary_text
            elif mode == 'supplychain':
                sess['consolidated_summary_supplychain'] = summary_text

        sess['last_analysis_mode'] = mode
        sess['active_file'] = None
        sess['active_topic'] = None

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

    mode_to_display = sess['last_analysis_mode']
    if mode_to_display is None:
        st.info("Click a button above to analyze Tariffs/Trade, Sales, or Supply Chain topics.")
        return

    if mode_to_display == 'tariff':
        processed_files = sess['processed_files_tariff']
        consolidated_summary = sess['consolidated_summary_tariff']
        sidebar_header = "Tariff Topics"
        main_header = "🌟 Consolidated Tariff & Trade Summary"
        agent = st.session_state.agent_tariff
    elif mode_to_display == 'sales':
        processed_files = sess['processed_files_sales']
        consolidated_summary = sess['consolidated_summary_sales']
        sidebar_header = "Sales Topics"
        main_header = "🌟 Consolidated Sales & Consumer Behavior Summary"
        agent = st.session_state.agent_sales
    elif mode_to_display == 'supplychain':
        processed_files = sess['processed_files_supplychain']
        consolidated_summary = sess['consolidated_summary_supplychain']
        sidebar_header = "Supply Chain Topics"
        main_header = "🌟 Consolidated Supply Chain & Logistics Summary"
        agent = st.session_state.agent_supplychain

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

    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(
                    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']

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

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

        # --- NEW UPDATED LOGIC ---
        # Send only the content of selected document(s):
        if len(selected_files) == 1:
            filename = selected_files[0]
            files_context = sess['processed_files_tariff'][filename]['context']
        else:
            files_context = "\n\n".join([
                f"=== Document: {filename} ===\n{sess['processed_files_tariff'][filename]['context']}"
                for filename in selected_files
            ])

        response_container = st.chat_message("assistant")
        placeholder = response_container.empty()
        full_response = ""

        try:
            response_stream, q_type, analysis_approach = st.session_state.chat_agent.stream_answer(prompt, files_context)

            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
                        placeholder.markdown(full_response + "▌")

            placeholder.markdown(full_response)

            sess['chat_history'].append({
                "role": "assistant",
                "content": full_response,
                "timestamp": datetime.utcnow().isoformat(),
                "analysis_type": q_type,
                "process": f"Answer generated from {'1 document' if len(selected_files)==1 else f'{len(selected_files)} documents'}"
            })

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

    if st.button("🔄 Reset Session", type="secondary"):
        preserved_keys = ["client", "agent_tariff", "agent_sales", "agent_supplychain", "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_sales': {},
            'processed_files_supplychain': {},
            'active_file': None,
            'active_topic': None,
            'topic_details_cache': {},
            'chat_history': [],
            'consolidated_summary_tariff': None,
            'consolidated_summary_sales': None,
            'consolidated_summary_supplychain': None,
            'last_analysis_mode': None
        }
        st.experimental_rerun()

if __name__ == "__main__":
    main()

give me the code as for diffenet diffent agent in differne files and all different stuctures