# app.py
import re

import streamlit as st
from langchain_core.messages import AIMessage, HumanMessage

# Import the agent initializer from its new location
from agent.agent import initialize_agent

def extract_image_paths(text: str) -> list[str]:
    """
    Finds all image filenames (png/jpeg) in a block of text,
    whether in quotes or bare.
    """
    return re.findall(r"['\"]?([A-Za-z0-9_\-]+\.(?:png|jpg|jpeg))['\"]?", text)


st.set_page_config(page_title="🤖 Smart Financial Analytics Agent", layout="wide")
st.title("🤖 CFO Copilot")

# Initialize the agent
agent_executor = initialize_agent()

# Initialize chat history in session state
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display past messages
for message in st.session_state.messages:
    with st.chat_message(message["role"]):

        st.text(message["content"])
        if "image_path" in message and message["image_path"]:
            st.image(message["image_path"])

# Get user input
if prompt := st.chat_input("Ask a question about your financial data..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Generate and display assistant response
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            chat_history = [
                HumanMessage(content=msg["content"]) if msg["role"] == "user" else AIMessage(content=msg["content"])
                for msg in st.session_state.messages[:-1]
            ]

            response = agent_executor.invoke({
                "input": prompt,
                "chat_history": chat_history
            })

            output_text = response["output"]
            st.text(output_text)

            # Collect any image paths from intermediate steps & output
            image_paths = []

            # 1. From intermediate_steps (even if action.tool != 'plot_chart')
            for step in response.get("intermediate_steps", []):
                _, observation = step
                # observation might be a filename or a descriptive text
                if isinstance(observation, str):
                    image_paths += extract_image_paths(observation)

            # 2. From the assistant’s final output text
            image_paths += extract_image_paths(output_text)

            # 3. De-duplicate and display
            for path in dict.fromkeys(image_paths):  # preserves order, removes dups
                try:
                    st.image(path)
                    # also record for session state
                    image_path = path
                except Exception as e:
                    st.error(f"Failed to load image {path}: {e}")

            # Save session state
            st.session_state.messages.append({
                "role": "assistant",
                "content": output_text,
                "image_path": image_path if image_paths else None
            })