from typing import TypedDict, Any, Dict
from langgraph.graph import StateGraph, END

from agents.grammar_agent import GrammarAgent
from agents.style_agent import StyleAgent
from agents.clarity_agent import ClarityAgent
from agents.reviewer_agent import ReviewerAgent




# ---------------------------
# TextDoctor State Definition
# ---------------------------
class TextDoctorState(TypedDict, total=False):
    input_text: str
    grammar_result: Dict[str, Any]
    style_result: Dict[str, Any]
    clarity_result: Dict[str, Any]
    review_result: Dict[str, Any]


# ---------------------------
# Initialize Agents Once
# ---------------------------
grammar_agent = GrammarAgent()
style_agent = StyleAgent()
clarity_agent = ClarityAgent()
reviewer_agent = ReviewerAgent()


# ---------------------------
# Node Definitions
# ---------------------------
def grammar_node(state: TextDoctorState) -> TextDoctorState:
    state["grammar_result"] = grammar_agent.correct(state["input_text"])
    return state


def style_node(state: TextDoctorState) -> TextDoctorState:
    g = state["grammar_result"]
    state["style_result"] = style_agent.stylize(g["corrected"])
    return state


def clarity_node(state: TextDoctorState) -> TextDoctorState:
    s = state["style_result"]
    state["clarity_result"] = clarity_agent.clarify(s["styled"])
    return state


def reviewer_node(state: TextDoctorState) -> TextDoctorState:
    state["review_result"] = reviewer_agent.review(
        original=state["input_text"],
        grammar_result=state["grammar_result"],
        style_result=state["style_result"],
        clarity_result=state["clarity_result"],
    )
    return state


# ---------------------------
# Build LangGraph Pipeline
# ---------------------------
graph = StateGraph(TextDoctorState)

graph.add_node("grammar", grammar_node)
graph.add_node("style", style_node)
graph.add_node("clarity", clarity_node)
graph.add_node("reviewer", reviewer_node)

graph.set_entry_point("grammar")
graph.add_edge("grammar", "style")
graph.add_edge("style", "clarity")
graph.add_edge("clarity", "reviewer")
graph.add_edge("reviewer", END)

# The compiled pipeline ready for use in UI
textdoctor_app = graph.compile()