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

# Conditionally import Azure and Google SDKs
try:
    import openai
    from openai import AzureOpenAI
except ImportError:
    openai = None

try:
    import google.generativeai as genai
except ImportError:
    genai = None


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


def parse_topics(text: str):
    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


# ===== Gemini (Google) Client Wrapper =====
class GeminiClient:
    def __init__(self):
        api_key = os.getenv("GOOGLE_API_KEY")
        if not api_key:
            st.error("Set GOOGLE_API_KEY environment variable before running")
            st.stop()
        if genai is None:
            st.error("Please install google-generativeai package to use Gemini model")
            st.stop()
        genai.configure(api_key=api_key)
        self.model = genai.GenerativeModel("gemini-2.0-flash")

    def generate(self, prompt: str, user_message: str) -> str:
        parts = [
            {"text": f"System: {prompt}"},
            {"text": f"User: {user_message}"}
        ]
        response = self.model.generate_content({"parts": parts})
        return response.text


# ===== Azure OpenAI Client Wrapper =====
class AzureClient:
    def __init__(self):
        api_key = os.getenv("AZURE_API_KEY")
        azure_endpoint = os.getenv("AZURE_ENDPOINT")
        deployment = os.getenv("AZURE_DEPLOYMENT")
        if not (api_key and azure_endpoint and deployment):
            st.error("Make sure AZURE_API_KEY, AZURE_ENDPOINT, and AZURE_DEPLOYMENT environment variables are set")
            st.stop()
        if openai is None:
            st.error("Please install openai package to use Azure OpenAI")
            st.stop()
        self.client = AzureOpenAI(api_key=api_key, azure_endpoint=azure_endpoint, api_version="2023-12-01-preview")
        self.deployment = deployment

    def generate(self, messages, temperature=0, stream=False):
        return self.client.chat.completions.create(
            model=self.deployment,
            messages=messages,
            temperature=temperature,
            stream=stream,
        )


# ===== Base Agent =====
class BaseAgent:
    def __init__(self, client, is_azure):
        self.client = client
        self.is_azure = is_azure

    def analyze_document(self, prompt_system, prompt_user, text):
        if self.is_azure:
            messages = [
                {"role": "system", "content": prompt_system},
                {"role": "user", "content": prompt_user + "\n\n" + text},
            ]
            response = self.client.generate(messages=messages, temperature=0)
            raw_text = response.choices[0].message.content.strip()
        else:
            raw_text = self.client.generate(prompt_system, prompt_user + "\n\n" + text)
        return parse_topics(raw_text)

    def generate_consolidated_summary_from_topics(self, prompt_system, prompt_user_template, processed_files):
        topics_block = ""
        for filename, filedata in processed_files.items():
            topics_block += f"### Topics from document: {filename}\n"
            for topic in sorted(filedata["topics"], key=lambda x: x["score"], reverse=True):
                topics_block += f"- **{topic['name']}**: {topic['summary']} (Relevance: {topic['score']})\n"
            topics_block += "\n"

        prompt_user = prompt_user_template.format(topics_block=topics_block)

        if self.is_azure:
            messages = [
                {"role": "system", "content": prompt_system},
                {"role": "user", "content": prompt_user},
            ]
            response = self.client.generate(messages=messages, temperature=0)
            return response.choices[0].message.content.strip()
        else:
            return self.client.generate(prompt_system, prompt_user)

    def generate_topic_detail(self, prompt_system, prompt_user_template, document_text, topic):
        prompt_user = prompt_user_template.format(text=document_text, topic=topic)
        if self.is_azure:
            messages = [
                {"role": "system", "content": prompt_system},
                {"role": "user", "content": prompt_user},
            ]
            response = self.client.generate(messages=messages, temperature=0, stream=True)
            collected_text = ""
            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
                        yield collected_text + "▌"
            yield collected_text
        else:
            text_out = self.client.generate(prompt_system, prompt_user)
            yield text_out


# ===== TariffTradeAgent =====
class TariffTradeAgent(BaseAgent):
    prompt_sys = "Tariff and trade policy analyst for Macy's CEO"
    prompt_analyze = """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:"""

    prompt_summary = """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."""

    prompt_detail = """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."""

    def analyze_document(self, text):
        return super().analyze_document(self.prompt_sys, self.prompt_analyze, text)

    def generate_consolidated_summary_from_topics(self, processed_files):
        return super().generate_consolidated_summary_from_topics(self.prompt_sys, self.prompt_summary, processed_files)

    def generate_topic_detail(self, text, topic):
        return super().generate_topic_detail(self.prompt_sys, self.prompt_detail, text, topic)


# ===== SalesAgent =====
class SalesAgent(BaseAgent):
    prompt_sys = "Sales analyst for Macy's executive team"
    prompt_analyze = """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:"""

    prompt_summary = """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."""

    prompt_detail = """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.
"""

    def analyze_document(self, text):
        return super().analyze_document(self.prompt_sys, self.prompt_analyze, text)

    def generate_consolidated_summary_from_topics(self, processed_files):
        return super().generate_consolidated_summary_from_topics(self.prompt_sys, self.prompt_summary, processed_files)

    def generate_topic_detail(self, text, topic):
        return super().generate_topic_detail(self.prompt_sys, self.prompt_detail, text, topic)


# ===== SupplyChainAgent =====
class SupplyChainAgent(BaseAgent):
    prompt_sys = "Supply chain analyst for Macy's"
    prompt_analyze = """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:"""

    prompt_summary = """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."""

    prompt_detail = """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.
"""

    def analyze_document(self, text):
        return super().analyze_document(self.prompt_sys, self.prompt_analyze, text)

    def generate_consolidated_summary_from_topics(self, processed_files):
        return super().generate_consolidated_summary_from_topics(self.prompt_sys, self.prompt_summary, processed_files)

    def generate_topic_detail(self, text, topic):
        return super().generate_topic_detail(self.prompt_sys, self.prompt_detail, text, topic)


# ===== EnhancedChatAgent for Q&A =====

class EnhancedChatAgent:
    def __init__(self, client, is_azure):
        self.client = client
        self.is_azure = is_azure
        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}
"""
        if self.is_azure:
            messages = [
                {"role": "system", "content": "Expert question classifier for Macy's analysis"},
                {"role": "user", "content": prompt},
            ]
            response = self.client.generate(messages=messages, temperature=0)
            return response.choices[0].message.content.strip().lower()
        else:
            return self.client.generate("Expert question classifier for Macy's analysis", prompt).strip().lower()

    def generate_answer(self, question, files_context):
        q_type = self.classify_question(question)

        if q_type == 'greetings':
            greeting_prompt = f"""
You are an senior 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:
"""
            if self.is_azure:
                messages = [
                    {"role": "system", "content": "Assistant specialized in greeting responses."},
                    {"role": "user", "content": greeting_prompt}
                ]
                response = self.client.generate(messages=messages, temperature=0)
                return response.choices[0].message.content.strip(), q_type
            else:
                return self.client.generate("Assistant specialized in greeting responses.", greeting_prompt), q_type

        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 asks about a single document and multiple selected, answer from only that one.
7. If user asks about multiple documents, synthesize across all.
8. Greetings get polite greeting responses only.

Current question type: {q_type.upper()}

Question: {question}

Relevant documents:
{files_context}

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

If information conflicts between documents, note explicitly.
"""
        if self.is_azure:
            messages = [
                {"role": "system", "content": "Senior analyst synthesizing information from multiple documents"},
                {"role": "user", "content": base_prompt},
            ]
            response = self.client.generate(messages=messages, temperature=0)
            return response.choices[0].message.content.strip(), q_type
        else:
            return self.client.generate("Senior analyst synthesizing information from multiple documents", base_prompt), q_type


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

def main():
    st.set_page_config(page_title="Tariff & Trade Analyst Pro: Azure & Gemini", layout="wide")
    st.title("📄 Macy's Document Analysis & Chat Tool")

    # Mode selection (Azure or Gemini)
    provider_select = st.radio("Select AI Provider", ["Azure GPT", "Google Gemini"])

    # Session Initialization
    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 or st.session_state.get("provider") != provider_select:
        # Initialize clients and agents based on selection
        if provider_select == "Azure GPT":
            st.session_state.client = AzureClient()
            is_azure = True
        else:
            st.session_state.client = GeminiClient()
            is_azure = False

        # Setup agents
        st.session_state.agent_tariff = TariffTradeAgent(st.session_state.client, is_azure)
        st.session_state.agent_sales = SalesAgent(st.session_state.client, is_azure)
        st.session_state.agent_supplychain = SupplyChainAgent(st.session_state.client, is_azure)
        st.session_state.chat_agent = EnhancedChatAgent(st.session_state.client, is_azure)
        st.session_state.provider = provider_select

    sess = st.session_state.session

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

    # --- Analysis Buttons Section ---
    st.subheader("2. Analyze Documents")
    col1, col2, col3 = st.columns(3)
    analyze_tariff = col1.button("🔍 Analyze Tariffs & Trade")
    analyze_sales = col2.button("📈 Analyze Sales & Consumer Behavior")
    analyze_supplychain = col3.button("🚚 Analyze Supply Chain & Logistics")

    mode = None
    agent = None
    processed_files = None

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

        if analyze_tariff:
            mode = "tariff"
            agent = st.session_state.agent_tariff
            processed_files = sess['processed_files_tariff']
        elif analyze_sales:
            mode = "sales"
            agent = st.session_state.agent_sales
            processed_files = sess['processed_files_sales']
        elif analyze_supplychain:
            mode = "supplychain"
            agent = st.session_state.agent_supplychain
            processed_files = sess['processed_files_supplychain']

        with st.spinner(f"Analyzing {mode} documents..."):
            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 extract text for {filename}")
                    except Exception as e:
                        st.error(f"Error processing {filename}: {str(e)}")

        with st.spinner("Generating consolidated summary..."):
            summary = agent.generate_consolidated_summary_from_topics(processed_files)
            if mode == "tariff":
                sess['consolidated_summary_tariff'] = summary
            elif mode == "sales":
                sess['consolidated_summary_sales'] = summary
            elif mode == "supplychain":
                sess['consolidated_summary_supplychain'] = summary

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

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

    # --- Show consolidated summary and topics ---
    last_mode = sess['last_analysis_mode']
    if last_mode is None:
        st.info("Select an analysis type above to get started.")
        return

    if last_mode == "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 last_mode == "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
    else:  # 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 for {last_mode.capitalize()} yet.")

    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"{last_mode}_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 = (last_mode, sess['active_file'], sess['active_topic'])
        if key_cache in sess['topic_details_cache']:
            st.markdown(sess['topic_details_cache'][key_cache])
        else:
            placeholder = st.empty()
            collected_text = ""

            try:
                detail_stream = agent.generate_topic_detail(
                    agent.prompt_sys,
                    agent.prompt_detail,
                    processed_files[sess['active_file']]['context'],
                    sess['active_topic'],
                )
                for chunk in detail_stream:
                    collected_text = chunk
                    placeholder.markdown(collected_text)
                sess['topic_details_cache'][key_cache] = collected_text
            except Exception as e:
                st.error(f"Error showing topic detail: {str(e)}")

    else:
        st.info("Select a file and topic from the sidebar to view detailed analysis.")

    st.divider()

    # --- Advanced Chat: Tariff & Trade Policies ---
    st.subheader("💬 Advanced Document Analysis (Tariff & Trade)")
    tariff_files = list(sess['processed_files_tariff'].keys())
    selected_files = st.multiselect("Select documents for question answering:", tariff_files)

    for msg in sess['chat_history']:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])

    question = st.chat_input("Ask a question about tariffs and trade policies")
    if question:
        if not selected_files:
            st.error("Select one or more documents for context before asking a question.")
            st.stop()

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

        if len(selected_files) == 1:
            context = sess['processed_files_tariff'][selected_files[0]]['context']
        else:
            context = "\n\n".join([
                f"=== Document: {fname} ===\n{sess['processed_files_tariff'][fname]['context']}"
                for fname in selected_files
            ])

        try:
            answer, q_type = st.session_state.chat_agent.generate_answer(question, context)
            sess['chat_history'].append({
                "role": "assistant",
                "content": answer,
                "timestamp": datetime.utcnow().isoformat(),
                "analysis_type": q_type
            })
            st.chat_message("assistant").markdown(answer)
        except Exception as e:
            st.error(f"Error during chat: {str(e)}")
            sess['chat_history'].append({
                "role": "assistant",
                "content": f"Error: {str(e)}",
                "timestamp": datetime.utcnow().isoformat(),
            })

    if st.button("🔄 Reset Session", type="secondary"):
        preserved_keys = ["client", "agent_tariff", "agent_sales", "agent_supplychain", "chat_agent", "provider"]
        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()


with this code i need to a folder structure of this appplication like agents are in folder , main class, utils, llm models, etc

give me full code