from langgraph.graph import StateGraph, END
from langchain_core.runnables import RunnableLambda
from src.agents import (
    CallState,
    IntakeAgent,
    TranscriptionAgent,
    SummarizationAgent,
    QualityScoringAgent,
    ModerationAgent,
    Router
)
from src.agents.Router import PostSummarizeRouter

def build_workflow():
    workflow = StateGraph(CallState)
    
    # Initialize agents
    intake_agent = IntakeAgent()
    transcription_agent = TranscriptionAgent()
    moderation_agent = ModerationAgent()
    summarization_agent = SummarizationAgent()
    quality_scoring_agent = QualityScoringAgent()
    router = Router()
    post_summarize_router = PostSummarizeRouter()
    
    # Add nodes
    workflow.add_node("intake", intake_agent)
    workflow.add_node("transcribe", transcription_agent)
    workflow.add_node("moderate", moderation_agent)
    summarize_node = RunnableLambda(summarization_agent).with_fallbacks(
        [RunnableLambda(SummarizationAgent.fallback)]
    )
    score_node = RunnableLambda(quality_scoring_agent).with_fallbacks(
        [RunnableLambda(QualityScoringAgent.fallback)]
    )
    workflow.add_node("summarize", summarize_node)
    workflow.add_node("score", score_node)
    
    # Define entry point
    workflow.set_entry_point("intake")
    
    # Add conditional edges from intake
    workflow.add_conditional_edges(
        "intake",
        router,
        {
            "transcribe": "transcribe",
            "summarize": "moderate",
            "end": END
        }
    )
    
    # Add standard edges
    workflow.add_edge("transcribe", "moderate")
    workflow.add_edge("moderate", "summarize")
    workflow.add_conditional_edges(
        "summarize",
        post_summarize_router,
        {"score": "score", "end": END},
    )
    workflow.add_edge("score", END)

    # Compile with a MemorySaver checkpointer when available.
    try:
        from langgraph.checkpoint.memory import MemorySaver

        return workflow.compile(checkpointer=MemorySaver())
    except Exception:
        return workflow.compile()