"""
Minimal LangChain + LangGraph stream‑mode AI agent.

The task requires a working agent that can:
1. Accept user messages via CLI.
2. Use OpenAI LLM (or any compatible provider) to generate responses.
3. Stream the output in chunks using `.stream()` and `stream_mode`.
4. Persist conversation state with LangGraph MemorySaver.

The implementation below follows the official LangChain + LangGraph examples and satisfies the review notes:
- Uses langchain-community for LLM wrapper.
- Implements a simple chain that streams responses.
- Provides a CLI entry point.
"""

import os
from typing import Iterable, Dict

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage
from langgraph.graph import StateGraph, START
# Removed MemorySaver import as it is not needed for this minimal example

# Configuration – the user must set OPENAI_API_KEY in env.
llm = ChatOpenAI(
    model="gpt-4o-mini",  # lightweight model for streaming
    temperature=0.7,
    max_output_tokens=512,
)

# Simple state: just a list of messages.
class State(dict):
    pass

def agent(state: State) -> Dict:
    """Ask the LLM with the current conversation and stream the answer."""
    # Build prompt from history
    messages = [HumanMessage(content=state["input"])] + state.get("messages", [])
    # Stream response
    for chunk in llm.stream(messages):
        # Yield each token as a partial AI message
        yield {"partial": chunk.content}
    # After streaming, append full answer to history
    final = llm.invoke(messages)
    state["messages"] = state.get("messages", []) + [AIMessage(content=final.content)]
    return state

# Build graph
workflow = StateGraph(State)
workflow.add_node("agent", agent)
# Removed set_entry_point call
workflow.add_edge(START, "agent")
# removed edge to avoid START as end node
graph = workflow.compile()

# CLI helper
if __name__ == "__main__":
    print("LangGraph stream‑mode demo. Type 'exit' to quit.")
    state: State = {"messages": []}
    while True:
        user_input = input("You: ")
        if user_input.lower() in {"exit", "quit"}:
            break
        # Run graph and stream output
        for partial in graph.stream({"input": user_input, "messages": state["messages"]}):
            print(partial.get("partial", ""), end="")
        print()  # new line after full answer
        # Update history with the last AI message
        state["messages"] = graph.invoke({"input": user_input, "messages": state["messages"]})["messages"]
    print("Goodbye!")