from smolagents import CodeAgent, HfApiModel, load_tool, tool, LiteLLMModel
import os
import datetime
import pytz
import yaml
from dotenv import load_dotenv
from tools.final_answer import FinalAnswerTool
from tools.query_knowledge import QueryKnowledgeTool
from tools.web_search import DuckDuckGoSearchTool
from tools.visit_webpage import VisitWebpageTool
from market_data.tools import (
    calculate_option_greeks,
    query_market_asset,
    query_option_chain,
    query_option_expirations,
    query_price_history,
    query_realized_volatility,
    query_volatility_snapshot,
)
from strategy.tools import build_volatility_strategy
from backtest.tools import analyze_strategy_payoff, backtest_long_straddle_csv, backtest_volatility_signal
from optimizer.tools import optimize_volatility_signal_parameters
from Gradio_UI import GradioUI

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    try:
        # Create timezone object
        tz = pytz.timezone(timezone)
        # Get current time in that timezone
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"


if __name__ == "__main__":
    final_answer = FinalAnswerTool()
    query_knowledge = QueryKnowledgeTool()
    web_search = DuckDuckGoSearchTool(max_results=6)
    visit_webpage = VisitWebpageTool()
    load_dotenv()
    hf_token = os.getenv("HF_TOKEN")
    gemini_api_key = os.getenv("GEMINI_API_KEY");
    
    model = LiteLLMModel(
        model_id="gemini/gemini-2.5-flash",
        temperature=0.2
    )
    
    with open("prompts.yaml", 'r') as stream:
        prompt_templates = yaml.safe_load(stream)
        
    agent = CodeAgent(
        model=model,
        tools=[
            query_market_asset,
            query_price_history,
            query_realized_volatility,
            query_option_expirations,
            query_option_chain,
            query_volatility_snapshot,
            calculate_option_greeks,
            build_volatility_strategy,
            analyze_strategy_payoff,
            backtest_long_straddle_csv,
            backtest_volatility_signal,
            optimize_volatility_signal_parameters,
            get_current_time_in_timezone,
            query_knowledge,
            web_search,
            visit_webpage,
            final_answer,
        ],
        max_steps=6,
        verbosity_level=1,
        grammar=None,
        planning_interval=None,
        name=None,
        description=None,
        prompt_templates=prompt_templates
    )
    
    
    GradioUI(agent).launch()