import streamlit as st
import yfinance as yf
from crewai import Agent, Task, Crew, Process, LLM
from crewai_tools import CodeInterpreterTool
from pydantic import BaseModel, Field
from dotenv import load_dotenv
import os

# Load environment variables
load_dotenv()

# Streamlit UI
st.title("AI-Powered Stock Analysis")

# Get Sambanova API Key from User
sambanova_key = st.text_input("Enter your Sambanova API Key", type="password")
query = st.text_input("Enter a stock query (e.g., 'Plot YTD stock gain of Tesla')")

# Define Query Output Model
class QueryAnalysisOutput(BaseModel):
    symbol: str = Field(..., description="Stock ticker symbol")
    timeframe: str = Field(..., description="Time period (e.g., '1d', '1mo', '1y')")
    action: str = Field(..., description="Action to be performed (e.g., 'fetch', 'plot')")

# Define LLM Model
if sambanova_key:
    llm = LLM(model="sambanova/DeepSeek-R1-Distill-Llama-70B", temperature=0.7, api_key=sambanova_key)
    
    # Define CrewAI Agents with backstory
    query_parser_agent = Agent(
        role="Stock Data Analyst",
        goal="Extract stock details from user query.",
        backstory="An expert in analyzing stock data and trends.",
        llm=llm,
        verbose=True,
        memory=True,
    )

    query_parsing_task = Task(
        description="Analyze the user query and extract stock details.",
        output_pydantic=QueryAnalysisOutput,
        agent=query_parser_agent,
    )

    code_writer_agent = Agent(
        role="Senior Python Developer",
        goal="Write Python code to visualize stock data.",
        backstory="A seasoned developer with expertise in financial data visualization.",
        llm=llm,
        verbose=True,
    )

    code_writer_task = Task(
        description="Generate Python code for stock visualization.",
        agent=code_writer_agent,
    )

    code_interpreter_tool = CodeInterpreterTool()
    code_execution_agent = Agent(
        role="Code Execution Expert",
        goal="Execute the generated code to visualize stock data.",
        backstory="Specializes in running and debugging Python scripts for data analysis.",
        tools=[code_interpreter_tool],
        allow_code_execution=True,
        llm=llm,
        verbose=True,
    )

    code_execution_task = Task(
        description="Execute the generated Python script.",
        agent=code_execution_agent,
    )

    # Create Crew
    crew = Crew(
        agents=[query_parser_agent, code_writer_agent, code_execution_agent],
        tasks=[query_parsing_task, code_writer_task, code_execution_task],
        process=Process.sequential,
    )

    if st.button("Analyze Stock"):
        result = crew.kickoff(inputs={"query": query})
        st.write("Results:", result)
else:
    st.warning("Please enter your Sambanova API Key to proceed.")