import os
import streamlit as st
from dotenv import load_dotenv
from phi import Agent
from phi.model.groq import Groq
from phi.tools.duckduckgo import DuckDuckGo
from phi.tools.yfinance import YFinanceTools

# Load environment variables
load_dotenv()

# Retrieve API keys from the environment
deepseek_api_key = os.getenv("GROQ_DEEPSEEK_API_KEY")
qwen_api_key = os.getenv("GROQ_QWEN_API_KEY")

# Debugging API key loading
if not deepseek_api_key or not qwen_api_key:
    st.error("Missing API keys. Please set them in Hugging Face Secrets.")
    st.stop()

# Define Web Agent
web_agent = Agent(
    name="Web Agent",
    model=Groq(id="qwen-2.5-coder-32b", api_key=qwen_api_key),
    tools=[DuckDuckGo()],
    instructions=["Always include sources"],
    show_tool_calls=True,
    markdown=True,
)

# Define Finance Agent
finance_agent = Agent(
    name="Finance Agent",
    role="Get financial data",
    model=Groq(id="qwen-2.5-coder-32b", api_key=qwen_api_key),
    tools=[YFinanceTools(stock_price=True, analyst_recommendations=True, company_info=True)],
    instructions=["Use tables to display data"],
    show_tool_calls=True,
    markdown=True,
)

# Create agent team
agent_team = Agent(
    model=Groq(id="deepseek-r1-distill-llama-70b", api_key=deepseek_api_key),
    team=[web_agent, finance_agent],
    instructions=["Always include sources", "Use tables to display data"],
    show_tool_calls=True,
    markdown=True,
)

# Streamlit UI
st.set_page_config(page_title="AI Chat Assistant", layout="wide")
st.title("🤖 AI Chat Assistant")

# User input
if prompt := st.chat_input("Ask me anything..."):
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        response_container = st.empty()
        response_text = ""

        # Check available response method
        if hasattr(agent_team, "respond") and callable(agent_team.respond):
            try:
                response_text = agent_team.respond(prompt) or "No response received."
                response_container.markdown(response_text)
            except Exception as e:
                error_message = f"Error: {str(e)}"
                st.error(error_message)
                response_text = error_message
        else:
            error_message = "Error: Agent does not support responses."
            st.error(error_message)
            response_text = error_message