from langgraph.graph import StateGraph, END
from state import KrishiState

from nodes.router import router_node
from nodes.disease import disease_node
from nodes.pesticide import pesticide_node
from nodes.weather import weather_node
from nodes.soil import soil_node
from nodes.crop import crop_node
from nodes.fertilizer import fertilizer_node
from nodes.advisor import advisor_node


builder = StateGraph(KrishiState)


builder.add_node("router", router_node)
builder.add_node("advisor", advisor_node)

builder.add_node("disease", disease_node)
builder.add_node("pesticide", pesticide_node)
builder.add_node("weather_d", weather_node)   # weather for disease branch
builder.add_node("soil_d", soil_node)         # soil for disease branch


builder.add_node("weather_i", weather_node)   # weather for irrigation branch
builder.add_node("soil_i", soil_node)         # soil for irrigation branch



builder.add_node("crop", crop_node)
builder.add_node("soil_f", soil_node)         # soil for fertilizer branch
builder.add_node("weather_f", weather_node)   # weather for fertilizer branch
builder.add_node("fertilizer", fertilizer_node)


builder.set_entry_point("router")

builder.add_conditional_edges(
    "router",
    lambda state: state["route"],
    {
        "disease":    "disease",
        "irrigation": "weather_i",
        "fertilizer": "crop",
    }
)

builder.add_edge("disease",   "pesticide")
builder.add_edge("pesticide", "weather_d")
builder.add_edge("weather_d", "soil_d")
builder.add_edge("soil_d",    "advisor")


builder.add_edge("weather_i", "soil_i")
builder.add_edge("soil_i",    "advisor")


builder.add_edge("crop",       "soil_f")
builder.add_edge("soil_f",    "weather_f")
builder.add_edge("weather_f",     "fertilizer")
builder.add_edge("fertilizer", "advisor")



#  TERMINAL

builder.add_edge("advisor", END)

graph = builder.compile()