import base64
from typing import List, TypedDict, Annotated, Optional
from langchain_openai import ChatOpenAI
from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage
from langgraph.graph.message import add_messages
from langgraph.graph import START, END, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition
from IPython.display import Image, display
from dotenv import load_dotenv

load_dotenv()  # reads OPENAI_API_KEY from .env

# How state got updated? (by Node return partial State update)
# In LangGraph, a Node typically does not mutate the shared State directly.
# It returns a partial State update, and LangGraph applies that update
# Node receives State → performs work → returns partial updates → LangGraph merges updates into State.
# lang chain llm message returned like this -> its pre defined
# AIMessage(
#     content="Hello! How can I help you?",
#     tool_calls=[],
#     id="run-123",
#     response_metadata={...},
#     usage_metadata={...}
# )

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

vision_llm = ChatOpenAI(model="gpt-4o") # model supports reading iages

def extract_text(img_path: str) -> str:
    """Extract text from an image using a vision model."""
    with open(img_path, "rb") as image_file:
        image_bytes = image_file.read()
    image_base64 = base64.b64encode(image_bytes).decode("utf-8") # base64 encode the image
    message = [
        HumanMessage(
            content=[
                {
                    "type": "text",
                    "text": "Extract all the text from this image. Return only the extracted text."
                },
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/png;base64,{image_base64}"
                    }
                }
            ]
        )
    ]

    response = vision_llm.invoke(message)
    return response.content.strip()


def divide(a: int, b: int) -> float:
    """Divide a and b."""
    return a / b

tools = [divide, extract_text]

llm = ChatOpenAI(model="gpt-4o")

llm_with_tools = llm.bind_tools(
    tools,
    parallel_tool_calls=False
)

#  بتستدعي الـ LLM وبترجع جوابه أو طلب استخدام أداة.
def assistant(state: AgentState):
    image = state["input_file"]
    sys_msg = SystemMessage(
        content=f"""
        You are a helpful butler named Alfred
        that serves Mr. Wayne and Batman.

        You can analyse documents and run computations
        with provided tools:

        extract_text(img_path: str) -> str:
            Extract text from an image file.

        divide(a: int, b: int) -> float:
            Divide a and b.

        Currently the loaded image is: {image}
        """
    )

    return {
        "messages": [
            llm_with_tools.invoke( # return AIMessage
                [sys_msg] + state["messages"]
            )
        ],
        "input_file": state["input_file"]
    }

# connect the graph:
# 1. Create the graph builder
builder = StateGraph(AgentState)
# 2. Add nodes
# The assistant node is responsible for both calling tools and generating the final response.
builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))
# 3. Set the starting point
builder.add_edge(START, "assistant")
# 4. Route based on tool calls work the same as tools_condition
def route_assistant(state: AgentState):
    last_message = state["messages"][-1]
    if last_message.tool_calls:
        return "tools"
    return END

builder.add_conditional_edges(
    "assistant",
    route_assistant,
    {
        "tools": "tools",
        END: END
    }
)
# 5. Return from tools to assistant
builder.add_edge("tools", "assistant")
# 6. Compile
react_graph = builder.compile()

############ call example: ############
initial_state = {
    "input_file": None,
    "messages": [
        HumanMessage(content="Divide 100 by 5")
    ]
}

#invoke() — executes the Graph and returns the final State.
#stream() — lets us observe what's happening during execution.

result = react_graph.invoke(initial_state)
print(result)