import os
from typing import TypedDict
from langgraph.graph import StateGraph, END
from oculus_sdk import OculusClient
import httpx

oculus = OculusClient(
    client_id=os.environ["AGENT_CLIENT_ID"],
    client_secret=os.environ["AGENT_CLIENT_SECRET"],
)

class AgentState(TypedDict):
    query: str
    data: dict | None
    result: str | None

def fetch_data_node(state: AgentState) -> AgentState:
    """Fetch data using a scoped Anti AI token."""
    token = oculus.get_token(scope="db:read")
    resp = httpx.get(
        "https://api.myservice.com/data",
        params={"q": state["query"]},
        headers={"Authorization": f"Bearer {token}"},
    )
    return {**state, "data": resp.json()}

def analyze_node(state: AgentState) -> AgentState:
    result = f"Analyzed {len(state['data'])} records"
    return {**state, "result": result}

graph = StateGraph(AgentState)
graph.add_node("fetch", fetch_data_node)
graph.add_node("analyze", analyze_node)
graph.add_edge("fetch", "analyze")
graph.add_edge("analyze", END)
graph.set_entry_point("fetch")

app = graph.compile()
result = app.invoke({"query": "top customers", "data": None, "result": None})
print(result["result"])