"""
Example: LangGraph checkpointer + long-term memory (recall) using AIEL SDK.

Required environment variables:
    AIEL_BASE_URL
    AIEL_TOKEN
    AIEL_WORKSPACE_ID
    AIEL_PROJECT_ID
    AIEL_OPENAI_CONNECTION_ID

Optional:
    AIEL_MODEL (default: gpt-4o-mini)
    AIEL_THREAD_ID (default: "demo-thread-1")
    AIEL_USER_ID (default: "u1")
"""

from __future__ import annotations

import os
from typing import Any, Dict, List, TypedDict

from aiel_sdk import AielClient
from aiel_sdk.adapters.langchain import AIELLangChainChatModel
from aiel_sdk.adapters.langgraph import AIELLangGraphMemory
from aiel_sdk.integrations.openai import ChatOpenAI, OpenAIChatModelAdapter
from aiel_sdk.memory import RemoteSaver

from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
from langgraph.store.base import BaseStore
from langgraph.store.base.batch import GetOp, PutOp
from langgraph.graph import END, START, StateGraph


def _env(name: str) -> str:
    value = (os.getenv(name) or "").strip()
    if not value:
        raise RuntimeError(f"Missing required environment variable: {name}")
    return value


class GraphState(TypedDict):
    messages: List[BaseMessage]
    user_id: str


def main() -> None:
    base_url = _env("AIEL_BASE_URL")
    token = _env("AIEL_TOKEN")
    workspace_id = _env("AIEL_WORKSPACE_ID")
    project_id = _env("AIEL_PROJECT_ID")
    connection_id = _env("AIEL_OPENAI_CONNECTION_ID")
    model_name = os.getenv("AIEL_MODEL", "gpt-4o-mini")
    thread_id = os.getenv("AIEL_THREAD_ID", "demo-thread-1")
    user_id = os.getenv("AIEL_USER_ID", "u1")

    client = AielClient(base_url=base_url, api_key=token)
    saver = RemoteSaver(client, workspace_id=workspace_id, project_id=project_id)
    memory = AIELLangGraphMemory(saver)
    checkpointer = memory.checkpointer
    store = memory.store

    chat = ChatOpenAI(
        client=client,
        workspace_id=workspace_id,
        project_id=project_id,
        connection_id=connection_id,
        model=model_name,
        temperature=0.2,
    )
    model = AIELLangChainChatModel(OpenAIChatModelAdapter(chat))

    def _store_get_profile(store_: BaseStore, user_id_: str) -> Dict[str, Any]:
        ns = ("users", user_id_)
        found = store_.batch([GetOp(namespace=ns, key="profile")])[0]
        if not found:
            return {}
        return dict(found.value or {})

    def _store_put_profile_fact(store_: BaseStore, user_id_: str, fact: str) -> None:
        ns = ("users", user_id_)
        store_.batch(
            [
                PutOp(
                    namespace=ns,
                    key="profile",
                    value={"last_fact": fact},
                )
            ]
        )

    def respond_node(state: GraphState, config: Dict[str, Any], *, store: BaseStore) -> Dict[str, Any]:
        _ = config
        profile_text = _store_get_profile(store, state["user_id"])

        response = model.invoke(
            [
                SystemMessage(
                    content=(
                        "You are a support assistant. Use user profile memory when helpful. "
                        f"Known profile: {profile_text}"
                    )
                ),
                *state["messages"],
            ]
        )
        return {"messages": [*state["messages"], AIMessage(content=str(response.content))]}

    def learn_profile_node(state: GraphState, config: Dict[str, Any], *, store: BaseStore) -> Dict[str, Any]:
        _ = config
        last_human = ""
        for msg in reversed(state["messages"]):
            if isinstance(msg, HumanMessage):
                last_human = str(msg.content)
                break
        if "my name is" in last_human.lower():
            _store_put_profile_fact(store, state["user_id"], last_human)
        return {}

    graph = StateGraph(GraphState)
    graph.add_node("respond", respond_node)
    graph.add_node("learn_profile", learn_profile_node)
    graph.add_edge(START, "respond")
    graph.add_edge("respond", "learn_profile")
    graph.add_edge("learn_profile", END)
    app = graph.compile(checkpointer=checkpointer, store=store)

    config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": "chat"}}

    # Turn 1
    out1 = app.invoke(
        {"messages": [HumanMessage(content="Hi, my name is Alex and I prefer concise answers.")], "user_id": user_id},
        config=config,
    )
    print("Turn 1:", out1["messages"][-1].content)

    # Turn 2 on same thread_id automatically uses checkpointed thread state.
    out2 = app.invoke(
        {"messages": [HumanMessage(content="What do you remember about me?")], "user_id": user_id},
        config=config,
    )
    print("Turn 2:", out2["messages"][-1].content)

    profile = _store_get_profile(store, user_id)
    print("Long memory profile:", profile)


if __name__ == "__main__":
    main()
