"""
LangGraph workflow for the Email Assistant multi-agent system.
Orchestrates the flow between different agents using LangGraph.
"""
from typing import Dict, Any, Literal, Annotated
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages

from models.state import EmailState, EmailInput
from agents import (
    InputParsingAgent,
    IntentDetectionAgent,
    DraftWriterAgent,
    ToneStylistAgent,
    ReviewAgent,
    RouterAgent
)


class EmailAssistantWorkflow:
    """Orchestrates the multi-agent workflow for email generation."""
    
    def __init__(self):
        # Initialize agents
        self.input_parser = InputParsingAgent()
        self.intent_detector = IntentDetectionAgent()
        self.draft_writer = DraftWriterAgent()
        self.tone_stylist = ToneStylistAgent()
        self.reviewer = ReviewAgent()
        self.router = RouterAgent()
        
        # Build the graph
        self.graph = self._build_graph()
    
    def _build_graph(self) -> StateGraph:
        """Build the LangGraph workflow."""
        
        # Create the graph with the state schema
        workflow = StateGraph(dict)
        
        # Add nodes for each agent
        workflow.add_node("parse_input", self._parse_input_node)
        workflow.add_node("detect_intent", self._detect_intent_node)
        workflow.add_node("write_draft", self._write_draft_node)
        workflow.add_node("style_tone", self._style_tone_node)
        workflow.add_node("review_draft", self._review_draft_node)
        workflow.add_node("handle_retry", self._handle_retry_node)
        workflow.add_node("finalize", self._finalize_node)
        workflow.add_node("handle_error", self._handle_error_node)
        
        # Set the entry point
        workflow.set_entry_point("parse_input")
        
        # Add edges for the main flow
        workflow.add_edge("parse_input", "detect_intent")
        workflow.add_edge("detect_intent", "write_draft")
        workflow.add_edge("write_draft", "style_tone")
        workflow.add_edge("style_tone", "review_draft")
        
        # Add conditional edges after review
        workflow.add_conditional_edges(
            "review_draft",
            self._route_after_review,
            {
                "retry": "handle_retry",
                "complete": "finalize",
                "error": "handle_error"
            }
        )
        
        # Retry goes back to draft writing
        workflow.add_edge("handle_retry", "write_draft")
        
        # Terminal nodes
        workflow.add_edge("finalize", END)
        workflow.add_edge("handle_error", END)
        
        return workflow.compile()
    
    def _parse_input_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """Node for parsing user input."""
        try:
            return self.input_parser.process(state)
        except Exception as e:
            state["error_message"] = f"Input parsing failed: {str(e)}"
            return state
    
    def _detect_intent_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """Node for detecting email intent."""
        try:
            return self.intent_detector.process(state)
        except Exception as e:
            state["error_message"] = f"Intent detection failed: {str(e)}"
            return state
    
    def _write_draft_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """Node for writing the email draft."""
        try:
            return self.draft_writer.process(state)
        except Exception as e:
            state["error_message"] = f"Draft writing failed: {str(e)}"
            return state
    
    def _style_tone_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """Node for styling the email tone."""
        try:
            return self.tone_stylist.process(state)
        except Exception as e:
            state["error_message"] = f"Tone styling failed: {str(e)}"
            return state
    
    def _review_draft_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """Node for reviewing the email draft."""
        try:
            return self.reviewer.process(state)
        except Exception as e:
            state["error_message"] = f"Review failed: {str(e)}"
            return state
    
    def _handle_retry_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """Node for handling retry logic."""
        return self.router.increment_retry(state)
    
    def _finalize_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """Node for finalizing the workflow."""
        return self.router.mark_complete(state)
    
    def _handle_error_node(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """Node for handling errors."""
        return self.router.handle_error(state)
    
    def _route_after_review(self, state: Dict[str, Any]) -> Literal["retry", "complete", "error"]:
        """Routing function after review."""
        return self.router.route_after_review(state)
    
    def run(self, email_input: EmailInput) -> EmailState:
        """Run the workflow with the given email input."""
        # Create initial state
        initial_state = EmailState(email_input=email_input)
        state_dict = initial_state.model_dump()
        
        # Run the graph
        final_state = self.graph.invoke(state_dict)
        
        # Return as EmailState
        return EmailState(**final_state)
    
    def run_dict(self, email_input: EmailInput) -> Dict[str, Any]:
        """Run the workflow and return the result as a dictionary."""
        initial_state = EmailState(email_input=email_input)
        state_dict = initial_state.model_dump()
        
        return self.graph.invoke(state_dict)


def create_workflow() -> EmailAssistantWorkflow:
    """Factory function to create a new workflow instance."""
    return EmailAssistantWorkflow()