import os
import yfinance as yf
from dotenv import load_dotenv
import streamlit as st
from langchain.llms import OpenAI
from langchain.agents import initialize_agent, Tool
from langchain.agents.agent_types import AgentType

# Load environment variables
load_dotenv()
api_key = os.getenv("AIML_API_KEY")

# Define custom Yahoo Finance tool
def get_stock_info(query: str) -> str:
    try:
        symbol = query.strip().split()[0].upper()
        stock = yf.Ticker(symbol)
        info = stock.info
        return f"{symbol} stock info:\nPrice: {info.get('currentPrice')}\nMarket Cap: {info.get('marketCap')}\nPE Ratio: {info.get('trailingPE')}"
    except Exception as e:
        return f"Error fetching stock data: {e}"

finance_tool = Tool(
    name="Stock Info Tool",
    func=get_stock_info,
    description="Use this tool to get basic stock info. Start your query with a stock ticker (e.g., NFLX for Netflix)."
)

# Initialize LLM
llm = OpenAI(
    openai_api_key=api_key,
    openai_api_base="https://api.aimlapi.com/v1",
    temperature=0.7
)

# Create agent
agent = initialize_agent(
    tools=[finance_tool],
    llm=llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True
)

# Streamlit UI
st.title("📈 Finance Assistant")
user_query = st.text_input("Enter your financial question:", "NFLX stock performance")

# Ensure the prompt is passed as a string
if st.button("Analyze"):
    with st.spinner("Analyzing..."):
        # Make sure user_query is passed as a string
        if isinstance(user_query, str):
            result = agent.run(user_query)  # This will now be passed correctly as a string
            st.markdown(result)
        else:
            st.error("Invalid query format. Please enter a valid question.")