"""
Runnable example: AIEL memory + OpenAI integration with LangChain and LangGraph.

Usage:
    export AIEL_BASE_URL="https://api.example.com"
    export AIEL_TOKEN="..."
    export AIEL_WORKSPACE_ID="ws_123"
    export AIEL_PROJECT_ID="proj_456"
    export AIEL_OPENAI_CONNECTION_ID="conn_openai_abc"
    export AIEL_MODEL="gpt-4o-mini"
    python -m aiel_sdk.examples.langchain_langgraph_memory
"""

from __future__ import annotations

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

from aiel_sdk import AielClient
from aiel_sdk.adapters.langgraph import AIELLangGraphMemory
from aiel_sdk.core.types import Message
from aiel_sdk.integrations.openai import ChatOpenAI, OpenAIChatModelAdapter
from aiel_sdk.memory import RemoteSaver, ThreadMessage

try:
    from aiel_sdk.adapters.langchain import AIELLangChainChatModel
except Exception as exc:  # pragma: no cover
    raise RuntimeError(
        "Missing LangChain dependencies. Install with: pip install 'aiel-sdk[langchain]'"
    ) from exc


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


def _build_clients() -> tuple[AielClient, RemoteSaver, OpenAIChatModelAdapter, AIELLangChainChatModel]:
    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 = os.getenv("AIEL_MODEL", "gpt-4o-mini")

    client = AielClient(base_url=base_url, api_key=token)
    saver = RemoteSaver(client, workspace_id=workspace_id, project_id=project_id)

    chat = ChatOpenAI(
        client=client,
        workspace_id=workspace_id,
        project_id=project_id,
        connection_id=connection_id,
        model=model,
        temperature=0.2,
    )
    core_model = OpenAIChatModelAdapter(chat)
    langchain_model = AIELLangChainChatModel(core_model)
    return client, saver, core_model, langchain_model


def _to_core_messages(thread_messages: List[ThreadMessage]) -> List[Message]:
    role_map = {
        "human": "user",
        "ai": "assistant",
        "system": "system",
        "tool": "tool",
    }
    out: List[Message] = []
    for msg in thread_messages:
        out.append(Message(role=role_map.get(msg.role, "user"), content=msg.content))
    return out


def run_langchain_roundtrip(
    saver: RemoteSaver,
    core_model: OpenAIChatModelAdapter,
    thread_key: str,
    user_text: str,
) -> str:
    thread = saver.thread_read(thread_key=thread_key, limit=50)
    messages = _to_core_messages(thread.messages)
    messages.append(Message(role="user", content=user_text))

    response = core_model.invoke(messages)
    answer = str(response.message.content)

    saver.thread_write(
        thread_key=thread_key,
        messages=[
            ThreadMessage(role="human", content=user_text),
            ThreadMessage(role="ai", content=answer),
        ],
    )
    saver.state_patch(thread_key=thread_key, patch={"last_user_input": user_text, "last_answer": answer})
    saver.checkpoint_write(
        thread_key=thread_key,
        state={"last_user_input": user_text, "last_answer": answer},
        metadata={"stage": "langchain_roundtrip"},
    )
    return answer


def run_langgraph_roundtrip(
    saver: RemoteSaver,
    core_model: OpenAIChatModelAdapter,
    thread_key: str,
    user_text: str,
) -> str:
    try:
        from langgraph.graph import END, START, StateGraph
    except Exception as exc:  # pragma: no cover
        raise RuntimeError(
            "Missing LangGraph dependency. Install with: pip install langgraph"
        ) from exc

    class State(dict):
        pass

    def model_node(state: Dict[str, Any]) -> Dict[str, Any]:
        messages = state["messages"]
        response = core_model.invoke(messages)
        return {"messages": [*messages, response.message]}

    graph = StateGraph(dict)
    graph.add_node("model", model_node)
    graph.add_edge(START, "model")
    graph.add_edge("model", END)
    memory = AIELLangGraphMemory(saver)
    checkpointer = memory.checkpointer
    app = graph.compile(checkpointer=checkpointer)

    thread = saver.thread_read(thread_key=thread_key, limit=50)
    base_messages = _to_core_messages(thread.messages)
    input_state = {"messages": [*base_messages, Message(role="user", content=user_text)]}

    output_state = app.invoke(input_state, config={"configurable": {"thread_id": thread_key}})
    final_message = output_state["messages"][-1]
    answer = str(final_message.content)

    saver.thread_write(
        thread_key=thread_key,
        messages=[
            ThreadMessage(role="human", content=user_text),
            ThreadMessage(role="ai", content=answer),
        ],
    )
    saver.checkpoint_write(
        thread_key=thread_key,
        state={"messages_count": len(output_state["messages"])},
        metadata={"stage": "langgraph_roundtrip"},
    )
    return answer


def main() -> None:
    _, saver, core_model, langchain_model = _build_clients()
    _ = langchain_model  # Kept for direct LangChain composition in your app.

    thread_key = os.getenv("AIEL_THREAD_KEY", f"example:{uuid.uuid4().hex[:10]}")
    user_text = os.getenv("AIEL_USER_TEXT", "Summarize why memory and checkpoints are useful in one sentence.")

    print(f"Thread key: {thread_key}")
    print("Running LangChain-style roundtrip...")
    answer_1 = run_langchain_roundtrip(saver, core_model, thread_key, user_text)
    print("LangChain answer:", answer_1)

    print("\nRunning LangGraph roundtrip...")
    answer_2 = run_langgraph_roundtrip(saver, core_model, thread_key, "Now answer in exactly five words.")
    print("LangGraph answer:", answer_2)

    state = saver.state_read(thread_key=thread_key)
    print("\nState keys:", sorted((state.state or {}).keys()))

    recall_ns = ["examples", "langchain_langgraph", "demo"]
    saver.recall_put(namespace=recall_ns, key="last_thread_key", value={"thread_key": thread_key})
    recall_item = saver.recall_get(namespace=recall_ns, key="last_thread_key")
    print("Recall item:", recall_item)


if __name__ == "__main__":
    main()
