from smolagents import CodeAgent, HfApiModel, tool, load_tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI
from bs4 import BeautifulSoup
from dotenv import load_dotenv
import os
import sys
sys.stdout.reconfigure(encoding='utf-8')


# ---------------------------------------------------------------------------
# Define a custom tool for investment advice.
# This tool fetches the top 5 headlines from Yahoo Finance,
# does a simple sentiment analysis by counting positive and negative words,
# and then provides advice based on the provided risk tolerance.
# ---------------------------------------------------------------------------

# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def my_custom_tool(arg1:str, arg2:int)-> str: #it's import to specify the return type
    #Keep this format for the description / args / args description but feel free to modify the tool
    """A tool that does nothing yet 
    Args:
        arg1: the first argument
        arg2: the second argument
    """
    return "What magic will you build ?"

@tool
def investment_advice_tool(risk_tolerance: str) -> str:
    """
    Provides investment advice based on current market sentiment and the investor's risk tolerance.
    
    Args:
        risk_tolerance: A string representing risk tolerance ('low', 'medium', or 'high').
    """
    try:
        # Fetch market news from Yahoo Finance
        url = "https://finance.yahoo.com/"
        response = requests.get(url)
        if response.status_code != 200:
            return f"Error fetching market news: HTTP {response.status_code}"
        
        # Parse the top headlines using BeautifulSoup
        soup = BeautifulSoup(response.text, "html.parser")
        headlines = soup.find_all("h3", limit=5)
        news = [headline.get_text() for headline in headlines]
        
        # Perform a simple sentiment analysis based on keywords
        positive_words = ["rise", "gain", "growth", "bullish", "record high"]
        negative_words = ["fall", "loss", "drop", "bearish", "recession"]
        sentiment_score = 0
        for article in news:
            for word in positive_words:
                if word in article.lower():
                    sentiment_score += 1
            for word in negative_words:
                if word in article.lower():
                    sentiment_score -= 1
        sentiment = "Bullish" if sentiment_score > 0 else "Bearish" if sentiment_score < 0 else "Neutral"
        
        # Provide advice based on risk tolerance and market sentiment
        if risk_tolerance.lower() == "low":
            if sentiment == "Bullish":
                advice = "Invest in stable index funds or blue-chip stocks for steady growth."
            elif sentiment == "Bearish":
                advice = "Consider bonds or fixed-income securities to minimize risk."
            else:
                advice = "Maintain a diversified portfolio focused on stability."
        elif risk_tolerance.lower() == "medium":
            if sentiment == "Bullish":
                advice = "Consider a balanced mix of growth stocks and dividend-paying stocks."
            elif sentiment == "Bearish":
                advice = "Look for undervalued defensive stocks and ETFs."
            else:
                advice = "A balanced portfolio combining growth and value stocks would be wise."
        elif risk_tolerance.lower() == "high":
            if sentiment == "Bullish":
                advice = "Explore high-growth sectors like tech or emerging markets for potentially high returns."
            elif sentiment == "Bearish":
                advice = "Adopt a cautious approach, perhaps allocating only a small portion to speculative investments."
            else:
                advice = "A portfolio with a mix of high-risk and stable investments might suit you."
        else:
            advice = "Invalid risk tolerance. Please choose 'low', 'medium', or 'high'."
        
        return f"Market sentiment is {sentiment}. Based on your {risk_tolerance} risk tolerance, {advice}"
    except Exception as e:
        return f"Error in investment advice tool: {str(e)}"


final_answer = FinalAnswerTool()

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 

model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
token=os.getenv("HF_TOKEN"),
custom_role_conversions=None,
)

# ---------------------------------------------------------------------------
# Load prompt templates from the prompts.yaml file.
# This file guides how the agent constructs prompts for the LLM.
# ---------------------------------------------------------------------------

# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
# ---------------------------------------------------------------------------
# Create the Financial Advisor Agent.
# We add our investment advice tool to the agent's tools list.
# ---------------------------------------------------------------------------
agent = CodeAgent(
    model=model,
    tools=[final_answer, investment_advice_tool],  # Add your custom financial tool here
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="Financial Advisor Agent",
    description="An agent that provides investment suggestions based on current market sentiment and your risk tolerance.",
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()