from langgraph.graph import StateGraph, END
from .state import PRReviewState
from .nodes import (
    code_retriever_node,
    code_reviewer_node,
    feedback_formatter_node,
    post_code_review_node,
    update_review_body_based_on_human_input_node,
)

def create_graph():
    """Creates and compiles the LangGraph for the PR review process."""
    workflow = StateGraph(PRReviewState)

    # --- Node Definitions ---
    workflow.add_node("code_retriever", code_retriever_node)
    workflow.add_node("code_reviewer", code_reviewer_node)
    workflow.add_node("feedback_formatter", feedback_formatter_node)
    workflow.add_node("post_code_review", post_code_review_node)
    workflow.add_node("update_review_body_based_on_human_input_node", update_review_body_based_on_human_input_node)

    # --- Edge Definitions ---
    workflow.set_entry_point("code_retriever")
    workflow.add_edge("code_retriever", "code_reviewer")
    workflow.add_edge("code_reviewer", "feedback_formatter")
    workflow.add_edge("feedback_formatter", "post_code_review")
    workflow.add_edge("post_code_review", "update_review_body_based_on_human_input_node")
    workflow.add_edge("update_review_body_based_on_human_input_node", END)

    # The compilation with the checkpointer is now handled in app.py
    return workflow

# To make it runnable for debugging or direct invocation
if __name__ == "__main__":
    app = create_graph()
    # You can add debug invocation here, e.g.:
    # from .state import PRReviewState
    # initial_state = PRReviewState(repo_name="your/repo", pr_id=123, ...)
    # for event in app.stream(initial_state):
    #     print(event)