"""
Example: 5-agent multi-agent workflow with LangChain + LangGraph + AIEL Memory.

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_KEY (default: generated)
    AIEL_TASK (default: built-in task)
"""

from __future__ import annotations

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

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

from langchain_core.messages import HumanMessage, SystemMessage
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 WorkflowState(TypedDict, total=False):
    thread_key: str
    task: str
    plan: str
    research: str
    draft: str
    critique: str
    final: str


def _invoke_agent(model: AIELLangChainChatModel, *, role_prompt: str, input_text: str) -> str:
    response = model.invoke(
        [
            SystemMessage(content=role_prompt),
            HumanMessage(content=input_text),
        ]
    )
    return str(response.content).strip()


def _build_runtime() -> tuple[RemoteSaver, AIELLangChainChatModel, str]:
    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") or "gpt-4o-mini").strip()

    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_name,
        temperature=0.2,
    )
    lc_model = AIELLangChainChatModel(OpenAIChatModelAdapter(chat))
    return saver, lc_model, model_name


def main() -> None:
    saver, lc_model, model_name = _build_runtime()

    thread_key = (os.getenv("AIEL_THREAD_KEY") or f"multi-agent:{uuid.uuid4().hex[:10]}").strip()
    task = (
        os.getenv("AIEL_TASK")
        or "Design a rollout plan for launching an AI support assistant for a SaaS product in 30 days."
    ).strip()

    # Persist user task as initial memory.
    saver.thread_write(
        thread_key=thread_key,
        messages=[ThreadMessage(role="human", content=task, meta={"agent": "user"})],
    )

    def planner_node(state: WorkflowState) -> WorkflowState:
        plan = _invoke_agent(
            lc_model,
            role_prompt=(
                "You are Planner Agent. Break the task into clear phases, constraints, "
                "risks, and success criteria. Keep it concise and structured."
            ),
            input_text=state["task"],
        )
        saver.thread_write(
            thread_key=state["thread_key"],
            messages=[ThreadMessage(role="ai", content=plan, meta={"agent": "planner"})],
        )
        return {"plan": plan}

    def researcher_node(state: WorkflowState) -> WorkflowState:
        research = _invoke_agent(
            lc_model,
            role_prompt=(
                "You are Researcher Agent. Expand the plan with assumptions, "
                "dependencies, metrics, and operational requirements."
            ),
            input_text=f"Task:\n{state['task']}\n\nPlan:\n{state['plan']}",
        )
        saver.thread_write(
            thread_key=state["thread_key"],
            messages=[ThreadMessage(role="ai", content=research, meta={"agent": "researcher"})],
        )
        return {"research": research}

    def writer_node(state: WorkflowState) -> WorkflowState:
        draft = _invoke_agent(
            lc_model,
            role_prompt=(
                "You are Writer Agent. Produce a practical rollout document "
                "with timeline, owners, milestones, and checklist."
            ),
            input_text=f"Task:\n{state['task']}\n\nPlan:\n{state['plan']}\n\nResearch:\n{state['research']}",
        )
        saver.thread_write(
            thread_key=state["thread_key"],
            messages=[ThreadMessage(role="ai", content=draft, meta={"agent": "writer"})],
        )
        return {"draft": draft}

    def critic_node(state: WorkflowState) -> WorkflowState:
        critique = _invoke_agent(
            lc_model,
            role_prompt=(
                "You are Critic Agent. Find gaps, unrealistic assumptions, missing risks, "
                "and unclear acceptance criteria. Return actionable improvements."
            ),
            input_text=f"Draft:\n{state['draft']}",
        )
        saver.thread_write(
            thread_key=state["thread_key"],
            messages=[ThreadMessage(role="ai", content=critique, meta={"agent": "critic"})],
        )
        return {"critique": critique}

    def finalizer_node(state: WorkflowState) -> WorkflowState:
        final = _invoke_agent(
            lc_model,
            role_prompt=(
                "You are Finalizer Agent. Merge draft + critique into one final deliverable. "
                "Keep clear sections and concrete next actions."
            ),
            input_text=f"Draft:\n{state['draft']}\n\nCritique:\n{state['critique']}",
        )
        saver.thread_write(
            thread_key=state["thread_key"],
            messages=[ThreadMessage(role="ai", content=final, meta={"agent": "finalizer"})],
        )
        saver.state_patch(
            thread_key=state["thread_key"],
            patch={
                "model": model_name,
                "task": state["task"],
                "completed_agents": ["planner", "researcher", "writer", "critic", "finalizer"],
            },
        )
        saver.checkpoint_write(
            thread_key=state["thread_key"],
            state={
                "plan": state.get("plan", ""),
                "research": state.get("research", ""),
                "draft": state.get("draft", ""),
                "critique": state.get("critique", ""),
                "final": final,
            },
            metadata={"workflow": "langchain_langgraph_multi_agent", "agents": 5},
        )
        return {"final": final}

    graph = StateGraph(WorkflowState)
    graph.add_node("planner", planner_node)
    graph.add_node("researcher", researcher_node)
    graph.add_node("writer", writer_node)
    graph.add_node("critic", critic_node)
    graph.add_node("finalizer", finalizer_node)
    graph.add_edge(START, "planner")
    graph.add_edge("planner", "researcher")
    graph.add_edge("researcher", "writer")
    graph.add_edge("writer", "critic")
    graph.add_edge("critic", "finalizer")
    graph.add_edge("finalizer", END)

    memory = AIELLangGraphMemory(saver)
    checkpointer = memory.checkpointer

    app = graph.compile(checkpointer=checkpointer)
    result = app.invoke(
        {"thread_key": thread_key, "task": task},
        config={"configurable": {"thread_id": thread_key, "checkpoint_ns": "multi-agent"}},
    )

    print(f"Thread key: {thread_key}")
    print("\n=== FINAL OUTPUT ===\n")
    print(result.get("final", "No final output generated."))


if __name__ == "__main__":
    main()
