from states.state import AgentState
import os
# Import the load_dotenv function from the dotenv library
from dotenv import load_dotenv
from langchain_google_genai import ChatGoogleGenerativeAI
from tools.multimodal_tools import extract_text, analyze_image_tool, analyze_audio_tool
from tools.math_tools import add, subtract, multiply, divide
from tools.search_tools import search_tool, serpapi_search
from tools.youtube_tools import extract_youtube_transcript
from langfuse.callback import CallbackHandler

load_dotenv()

# Read your API key from the environment variable or set it manually
api_key = os.getenv("GEMINI_API_KEY")
langfuse_secret_key = os.getenv("LANGFUSE_SECRET_KEY")
langfuse_public_key = os.getenv("LANGFUSE_PUBLIC_KEY")

# Initialize Langfuse CallbackHandler for LangGraph/Langchain (tracing)
langfuse_handler = CallbackHandler(
    public_key=langfuse_public_key,
    secret_key=langfuse_secret_key,
    host="http://localhost:3000"
)

chat = ChatGoogleGenerativeAI(
    model= "gemini-2.5-pro-preview-05-06",
    temperature=0,
    max_retries=2,
    google_api_key=api_key,
    thinking_budget= 0
)

tools = [
    extract_text,
    analyze_image_tool,
    analyze_audio_tool,
    extract_youtube_transcript,
    add,
    subtract,
    multiply,
    divide,
    search_tool
]

chat_with_tools = chat.bind_tools(tools)

def assistant(state: AgentState):
    sys_msg = "You are a helpful assistant with access to tools. Understand user requests accurately. Use your tools when needed to answer effectively. Strictly follow all user instructions and constraints." \
    "Pay attention: your output needs to contain only the final answer without any reasoning since it will be strictly evaluated against a dataset which contains only the specific response." \
    "Your final output needs to be just the string or integer containing the answer, not an array or technical stuff."
    return {
        "messages": [chat_with_tools.invoke([sys_msg] + state["messages"])]
    }