import streamlit as st
import pandas as pd
import os
import asyncio
from llama_index.llms.openai import OpenAI
from llama_index.core.agent import ReActAgent
from llama_index.core.workflow import Context 
from llama_index.core.tools import FunctionTool

# --- Import tools class ---
from tools import PokemonAdvisorTools

# --- 1. Page Configuration (MUST be first Streamlit command) ---
st.set_page_config(page_title="cAsh Robo-Advisor", page_icon="🤖")

# --- 2. Initialize Agent & Tools (with caching to avoid recreating on every rerun) ---
@st.cache_resource
def initialize_agent():
    """Initialize the agent and tools once and cache them."""
    # Instantiate the class to load data
    advisor = PokemonAdvisorTools("/app/src")
    
    # Create the list of tool methods to wrap
    tool_methods = [
        advisor.get_card_info,
        advisor.find_grading_opportunities,
        advisor.assess_risk_volatility,
        advisor.get_roi_metrics,
        advisor.get_recent_price_spikes,
        advisor.analyze_set_performance,
        advisor.find_cards_by_artist,
        advisor.get_market_movers
    ]
    
    # Wrap tools into LlamaIndex FunctionTools
    tools = [FunctionTool.from_defaults(fn=func) for func in tool_methods]
    
    # Initialize the LLM (Ensure OPENAI_API_KEY is set in your env)
    llm = OpenAI(model="gpt-4o-mini", temperature=0.5)
    
    # System Prompt
    system_prompt = """
        ### ROLE
        You are **cAsh**, a sophisticated Pokemon Investment Advisor. 
        While you rely **strictly** on data and tools, your goal is to provide a helpful, engaging experience similar to a professional financial consultant.

        ### CONVERSATIONAL STYLE
        - **Opening:** Start with a brief, friendly acknowledgment of the user's request.
        - **Interpretation:** Don't just list numbers. Explain *why* these numbers matter (e.g., "The market is showing a significant spike for high-end Tag Team cards...").
        - **Formatting:** Use bolding for card names and bullet points for readability, but keep the sentences flowing naturally.
        - **Closing:** End with a helpful follow-up suggestion or a brief summary of the risk/opportunity.

        ### TOOL USAGE PROTOCOL
        1. **Verify First:** ALWAYS use `get_card_info` first for specific card queries.
        2. **Safety Check:** BEFORE recommending ANY card, you MUST call `assess_risk_volatility`.
        3. **Data Integrity:** Never hallucinate. If the tool returns nothing, say "My database doesn't have the specific stats for that card yet."

        ### TONE
        Professional, objective, insightful, and conversational. Always answer in English.
        """
    
    # Initialize Agent
    agent = ReActAgent(
        tools=tools,
        llm=llm,
        verbose=True,
        system_prompt=system_prompt,
        timeout=120
    )
    
    return agent

# Get cached agent
agent = initialize_agent()

# --- 3. Initialize Session State ---
if "messages" not in st.session_state:
    st.session_state.messages = []

if "context" not in st.session_state:
    # Create a new Context for this session
    st.session_state.context = Context(agent)

# --- 4. Define the Chat Function (Sync wrapper for async) ---
async def ask_advisor_async(user_message, ctx):
    """Async function to handle the chat."""
    if not user_message:
        return "Please enter a message."
    
    try:
        # Execute the agent workflow
        response = await agent.run(user_msg=user_message, ctx=ctx)
        return str(response)
    except Exception as e:
        return f"⚠️ **Agent Error:** {str(e)}\n\n*Check the console logs for detailed tool output.*"

def ask_advisor(user_message):
    """Synchronous wrapper for the async function."""
    return asyncio.run(ask_advisor_async(user_message, st.session_state.context))

# --- 5. Styling & Header ---
st.title("🍃 🔥 cAsh 💧 ⚡")
st.markdown("""
**🤖 Your AI Quantitative Analyst for Pokemon Cards.**
"""
)

# --- 6. Sidebar Examples ---
st.sidebar.header("Example Queries")
examples = [
    "What are the top 3 grading opportunities right now?",
    "What cards are trending down?",
    "Show me profitable cards by Tomokazu Komiya."
]

for ex in examples:
    if st.sidebar.button(ex, key=f"example_{ex}"):
        # Add to messages and set a flag to process
        st.session_state.messages.append({"role": "user", "content": ex})
        # Process the example query
        with st.spinner("Thinking..."):
            response = ask_advisor(ex)
        st.session_state.messages.append({"role": "assistant", "content": response})
        st.rerun()

# --- 7. Display Chat History ---
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# --- 8. Chat Input Logic ---
if prompt := st.chat_input("Ask cAsh..."):
    # Display user message
    with st.chat_message("user"):
        st.markdown(prompt)
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # Generate Response
    with st.chat_message("assistant"):
        with st.spinner("Analyzing..."):
            response = ask_advisor(prompt)
        st.markdown(response)
    st.session_state.messages.append({"role": "assistant", "content": response})
    
    # Force a rerun to update the chat display
    st.rerun()