import os
from typing import TypedDict, Annotated
from dotenv import load_dotenv
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage
from langgraph.prebuilt import ToolNode
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import tools_condition
from langchain_anthropic import ChatAnthropic

from utils.read_file import read_txt_file
from tools import available_tools

SYS_PROMPT_PATH = "./sys_prompt.txt"

load_dotenv()
api_key = os.getenv("ANTHROPIC_API_KEY")

llm = ChatAnthropic(
    model="claude-3-7-sonnet-20250219",
    temperature=0.5,
    api_key=api_key,
    # max_tokens=5000,
    # thinking={"type": "enabled", "budget_tokens": 2000}
    )
tools = available_tools
llm_with_tools = llm.bind_tools(tools)

class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]

def assistant(state: AgentState):
    return {
        "messages": [llm_with_tools.invoke(state["messages"])],
    }

builder = StateGraph(AgentState)

builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))

builder.add_edge(START, "assistant")
builder.add_conditional_edges(
    "assistant",
    tools_condition
)
builder.add_edge("tools","assistant")

my_agent = builder.compile()

SYSTEM_PROMPT = read_txt_file(SYS_PROMPT_PATH)