"""
Einfaches LangGraph-Beispiel
"""

from typing import TypedDict, Optional
from langgraph.graph import StateGraph, END

# Dummy-LLM zum Testen
class DummyLLM:
    def invoke(self, text):
        if "hoch" in text.lower():
            return "hoch"
        elif "mittel" in text.lower():
            return "mittel"
        else:
            return "niedrig"

llm = DummyLLM()

# Zustandstyp definieren
class State(TypedDict):
    question: str
    risk_level: Optional[str]
    response: Optional[str]

# Graph erstellen
graph = StateGraph(State)

def classify(state):
    # Risiko klassifizieren
    question = state["question"]
    risk_level = llm.invoke(f"Klassifiziere das Risiko: {question}")
    return {"risk_level": risk_level}

def handle(state):
    # Risiko behandeln basierend auf Klassifizierung
    risk_level = state["risk_level"]
    if "hoch" in risk_level.lower():
        response = "Hohes Risiko: Eskalation an Senior Risk Manager"
    elif "mittel" in risk_level.lower():
        response = "Mittleres Risiko: Standard-Prüfprozess einleiten"
    else:
        response = "Niedriges Risiko: Automatische Freigabe"
    return {"response": response}

# Knoten hinzufügen
graph.add_node("classify", classify)
graph.add_node("handle", handle)

# Kanten hinzufügen
graph.add_edge("classify", "handle")
graph.add_edge("handle", END)

# Startknoten definieren
graph.set_entry_point("classify")

# Graph kompilieren
app = graph.compile()

# Graph testen
test_input = {"question": "Was tun bei hoher Marktvolatilität?"}
result = app.invoke(test_input)
print("Test result:", result)

# Optional: Graph visualisieren und als PNG speichern
try:
    from IPython.display import Image
    from langchain_core.runnables.graph import MermaidDrawMethod
    
    # Generate PNG using Mermaid API
    png_data = app.get_graph().draw_mermaid_png(
        draw_method=MermaidDrawMethod.API
    )
    
    # Save the PNG to a file
    with open("simple_graph.png", "wb") as f:
        f.write(png_data)
    print("\nGraph visualization saved to simple_graph.png")
    
except ImportError:
    print("\nVisualization packages not available.")
    print("To visualize the graph, install: pip install ipython pillow langchain-core")