graph.py

from typing import Literal
from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from state import ConversationState
from agents import (
    RouterAgent, 
    GeneralDesignAgent, 
    BudgetAnalysisAgent, 
    FloorplanAgent,
    FloorplanGeneratorAgent,
    RegulationAgent
)
from detailed_budget_agent import DetailedBudgetAgent
# Removed creative_specialists - they had NoneType formatting errors
# Keeping only the core working agents


def create_initial_state() -> ConversationState:
    """Create initial conversation state with enhanced multi-agent memory"""
    return {
        "messages": [],
        "current_topic": None,
        "user_requirements": {
            "budget": None,
            "location": "Montreal",  # Default as specified in requirements
            "family_size": None,
            "lifestyle_preferences": [],
            "special_needs": []
        },
        "floorplan_requirements": {
            "num_floors": None,
            "total_sqft": None,
            "lot_shape": None,
            "lot_dimensions": None,
            "rooms": []
        },
        "detailed_floorplan": {
            "design_analysis": None,
            "detailed_rooms": [],
            "structural_elements": [],
            "circulation_plan": {},
            "lot_utilization": {},
            "architectural_features": []
        },
        "budget_breakdown": {
            "site_preparation": None,
            "foundation": None,
            "framing": None,
            "roofing": None,
            "exterior_finishes": None,
            "interior_finishes": None,
            "mechanical_systems": None,
            "electrical_systems": None,
            "plumbing_systems": None,
            "permits_fees": None,
            "professional_services": None,
            "contingency": None,
            "total_construction_cost": None,
            "cost_per_sqft": None,
            "budget_analysis": None
        },
        "conversation_history": [],
        "agent_recommendations": [],
        "agent_memory": {},  # Shared memory between agents
        "next_agent": None,
        "floorplan_ready": False,
        "budget_ready": False
    }


def should_generate_floorplan(state: ConversationState) -> Literal["generate_floorplan", "route_to_agent"]:
    """Decide whether to generate floorplan or continue conversation"""
    if state["floorplan_ready"]:
        return "generate_floorplan"
    return "route_to_agent"


def should_generate_budget(state: ConversationState) -> Literal["generate_budget", "route_to_agent"]:
    """Decide whether to generate detailed budget or continue conversation"""
    # Check if we have enough info for detailed budget and floorplan is designed
    has_budget = state["user_requirements"]["budget"] is not None
    has_floorplan_details = state["detailed_floorplan"]["detailed_rooms"] != []
    
    if has_budget and has_floorplan_details and not state["budget_ready"]:
        return "generate_budget"
    return "route_to_agent"


def route_to_specialist(state: ConversationState) -> Literal[
    "general_design", "budget_analysis", "floorplan", "detailed_budget", "regulation", "end"
]:
    """Route to appropriate specialist agent - core working agents plus regulation"""
    next_agent = state.get("next_agent")
    
    # Check if we should do detailed budget after floorplan
    if state["detailed_floorplan"]["detailed_rooms"] and not state["budget_ready"]:
        return "detailed_budget"
    
    # Route to specialist based on agent decision - core working agents
    routing_map = {
        "general": "general_design",
        "budget": "budget_analysis", 
        "floorplan": "floorplan",
        "regulation": "regulation"
    }
    
    return routing_map.get(next_agent, "end")


def create_architecture_assistant_graph(model: ChatOpenAI) -> StateGraph:
    """Create the LangGraph workflow for the architecture assistant with multi-agent collaboration"""
    
    # Initialize core working agents plus regulation
    router = RouterAgent(model)
    general_agent = GeneralDesignAgent(model)
    budget_agent = BudgetAnalysisAgent(model)
    floorplan_agent = FloorplanAgent(model)
    floorplan_generator = FloorplanGeneratorAgent(model)
    detailed_budget_agent = DetailedBudgetAgent(model)
    regulation_agent = RegulationAgent(model)
    
    # Create the graph
    workflow = StateGraph(ConversationState)
    
    # Add core working nodes plus regulation
    workflow.add_node("router", router.process)
    workflow.add_node("general_design", general_agent.process)
    workflow.add_node("budget_analysis", budget_agent.process)
    workflow.add_node("floorplan", floorplan_agent.process)
    workflow.add_node("generate_floorplan", floorplan_generator.process)
    workflow.add_node("detailed_budget", detailed_budget_agent.process)
    workflow.add_node("regulation", regulation_agent.process)
    
    # Set entry point
    workflow.set_entry_point("router")
    
    # Add conditional routing from router to core working specialists plus regulation
    workflow.add_conditional_edges(
        "router",
        route_to_specialist,
        {
            "general_design": "general_design",
            "budget_analysis": "budget_analysis", 
            "floorplan": "floorplan",
            "detailed_budget": "detailed_budget",
            "regulation": "regulation",
            "end": END
        }
    )
    
    # Add conditional routing from floorplan agent to check if ready for generation
    workflow.add_conditional_edges(
        "floorplan",
        should_generate_floorplan,
        {
            "generate_floorplan": "generate_floorplan",
            "route_to_agent": END
        }
    )
    
    # After floorplan generation, automatically trigger detailed budget if budget available
    workflow.add_conditional_edges(
        "generate_floorplan",
        should_generate_budget,
        {
            "generate_budget": "detailed_budget",
            "route_to_agent": END
        }
    )
    
    # All core agents end conversation after responding (user can continue with new input)
    workflow.add_edge("general_design", END)
    workflow.add_edge("budget_analysis", END)
    workflow.add_edge("detailed_budget", END)
    workflow.add_edge("regulation", END)
    
    return workflow.compile()


class ArchitectureAssistant:
    """Main architecture assistant class with persistent state management"""
    
    def __init__(self, openai_api_key: str, user_id: str = None):
        self.model = ChatOpenAI(
            api_key=openai_api_key,
            model="gpt-4o-mini",
            temperature=0.7
        )
        self.graph = create_architecture_assistant_graph(self.model)
        self.state = create_initial_state()
        
        # Initialize user state management
        from user_state_manager import user_state_manager
        self.state_manager = user_state_manager
        self.user_id = user_id
        self.session_id = None
        
        # Start new session
        if user_id:
            self.session_id = self.state_manager.start_new_session(user_id)
    
    def chat(self, user_input: str, save_state: bool = True) -> str:
        """Process user input and return response with optional state saving"""
        # Add user message to state
        self.state["messages"].append({
            "role": "user",
            "content": user_input
        })
        
        # Process through the graph
        result = self.graph.invoke(self.state)
        
        # Update state with result
        self.state = result
        
        # Save state after each interaction
        if save_state and self.user_id:
            self.state_manager.save_user_state(
                self.state, 
                self.user_id, 
                self.session_id
            )
        
        # Return the last assistant message
        assistant_messages = [msg for msg in self.state["messages"] if msg["role"] == "assistant"]
        if assistant_messages:
            return assistant_messages[-1]["content"]
        else:
            return "I'm here to help with your home design questions!"
    
    def get_conversation_summary(self) -> dict:
        """Get a summary of the current conversation state"""
        return {
            "user_requirements": self.state["user_requirements"],
            "floorplan_requirements": self.state["floorplan_requirements"],
            "current_topic": self.state["current_topic"],
            "total_messages": len(self.state["messages"])
        }
    
    def reset_conversation(self, start_new_session: bool = True):
        """Reset the conversation state"""
        self.state = create_initial_state()
        
        # Start new session if requested
        if start_new_session and self.user_id:
            self.session_id = self.state_manager.start_new_session(self.user_id)
    
    def load_previous_state(self, session_id: str = None) -> bool:
        """Load a previous conversation state"""
        if not self.user_id:
            return False
        
        loaded_state = self.state_manager.load_user_state(self.user_id, session_id)
        if loaded_state:
            self.state = loaded_state
            if session_id:
                self.session_id = session_id
            return True
        
        return False
    
    def get_user_history(self) -> list:
        """Get conversation history for current user"""
        if not self.user_id:
            return []
        
        return self.state_manager.get_user_history(self.user_id)
    
    def set_user_id(self, user_id: str):
        """Set or change the user ID"""
        self.user_id = user_id
        self.session_id = self.state_manager.start_new_session(user_id)