from typing import Dict, Any
from langchain_openai import ChatOpenAI
from langchain.schema import HumanMessage, SystemMessage
from state import ConversationState


class BaseAgent:
    def __init__(self, model: ChatOpenAI):
        self.model = model
    
    def process(self, state: ConversationState) -> ConversationState:
        raise NotImplementedError


class RouterAgent(BaseAgent):
    """Intelligently routes user messages to appropriate specialist agents"""
    
    def process(self, state: ConversationState) -> ConversationState:
        user_message = state["messages"][-1]["content"] if state["messages"] else ""
        
        # Check project progress to suggest next logical step
        project_status = self._analyze_project_status(state)
        
        system_prompt = f"""You are an intelligent routing agent for a comprehensive residential architecture assistant. 
        
        Available specialists:
        - "general": General home design questions, architecture principles, getting started
        - "budget": Budget analysis, cost estimates, Montreal market reality checks
        - "floorplan": Room layout, square footage, floor planning, lot dimensions - CAN GENERATE DESIGNS WITH MINIMAL INFO
        - "regulation": Montreal building codes, permits, zoning requirements, regulatory compliance
        
        Current project status: {project_status}
        
        IMPORTANT: Our floorplan specialist can now create initial designs with very basic requirements!
        Route to "floorplan" if user mentions ANY room needs, house size, or design requests.
        Don't wait for complete requirements - we show designs first, then iterate.
        
        Route based on:
        1. User's explicit request
        2. Natural conversation flow 
        3. Proactive design generation (route to floorplan early and often)
        4. Focus on the three core working specialists only
        
        Respond with only the specialist name."""
        
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=f"User message: {user_message}")
        ]
        
        response = self.model.invoke(messages)
        next_agent = response.content.strip().lower()
        
        state["next_agent"] = next_agent
        state["current_topic"] = next_agent
        
        return state
    
    def _analyze_project_status(self, state: ConversationState) -> str:
        """Analyze current project status to guide routing decisions"""
        status_parts = []
        
        # Check what's been completed - only track core working features
        if state["detailed_floorplan"]["detailed_rooms"]:
            status_parts.append("architectural design complete")
        if state["budget_breakdown"]["total_construction_cost"]:
            status_parts.append("budget estimated")
        
        if not status_parts:
            return "project just starting, gathering basic requirements"
        elif len(status_parts) == 1:
            return f"early design phase - {', '.join(status_parts)}"
        else:
            return f"design complete - {', '.join(status_parts)} - ready for construction planning"


class GeneralDesignAgent(BaseAgent):
    """Handles general home design questions and architectural advice"""
    
    def process(self, state: ConversationState) -> ConversationState:
        user_message = state["messages"][-1]["content"]
        context = self._build_context(state)
        
        system_prompt = f"""You are an expert residential architect and home design consultant.
        
        Current context about the user:
        {context}
        
        Provide helpful, practical advice about home design, architecture, and building.
        Keep responses conversational and reference their context when relevant.
        If they seem ready to discuss budget or floorplan specifics, gently guide them there."""
        
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=user_message)
        ]
        
        response = self.model.invoke(messages)
        
        state["messages"].append({
            "role": "assistant", 
            "content": response.content,
            "agent": "general_design"
        })
        
        return state
    
    def _build_context(self, state: ConversationState) -> str:
        context_parts = []
        
        if state["user_requirements"]["budget"]:
            context_parts.append(f"Budget: ${state['user_requirements']['budget']:,.0f}")
        
        if state["user_requirements"]["location"]:
            context_parts.append(f"Location: {state['user_requirements']['location']}")
            
        if state["user_requirements"]["family_size"]:
            context_parts.append(f"Family size: {state['user_requirements']['family_size']}")
            
        if state["user_requirements"]["lifestyle_preferences"]:
            context_parts.append(f"Preferences: {', '.join(state['user_requirements']['lifestyle_preferences'])}")
        
        return "\n".join(context_parts) if context_parts else "No specific requirements collected yet."


class BudgetAnalysisAgent(BaseAgent):
    """Analyzes budgets and provides Montreal market reality checks"""
    
    def process(self, state: ConversationState) -> ConversationState:
        user_message = state["messages"][-1]["content"]
        context = self._build_context(state)
        
        system_prompt = f"""You are a construction cost analyst specializing in Montreal residential market.
        
        Current context about the user:
        {context}
        
        Montreal Market Knowledge:
        - New construction: $300-500 CAD per sq ft (varies by finishes)
        - Land costs: $200-400k+ depending on area
        - Permits and fees: $10-20k typical
        - Architecture/engineering: 8-15% of construction cost
        
        Help users understand realistic budgets, extract their budget if not known, and guide them toward 
        appropriate home sizes and features. Be encouraging but realistic about Montreal costs."""
        
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=user_message)
        ]
        
        response = self.model.invoke(messages)
        
        # Extract budget information if mentioned
        self._extract_budget_info(user_message, state)
        
        state["messages"].append({
            "role": "assistant", 
            "content": response.content,
            "agent": "budget_analysis"
        })
        
        return state
    
    def _extract_budget_info(self, message: str, state: ConversationState):
        """Extract budget and family info from user message"""
        import re
        
        # Look for budget mentions
        budget_patterns = [
            r'\$?(\d{1,3}(?:,\d{3})*(?:\.\d{2})?)\s*(?:k|thousand)',
            r'\$?(\d{1,3}(?:,\d{3})*(?:\.\d{2})?)',
        ]
        
        for pattern in budget_patterns:
            matches = re.findall(pattern, message.lower())
            if matches:
                try:
                    budget = float(matches[0].replace(',', ''))
                    if 'k' in message.lower() or 'thousand' in message.lower():
                        budget *= 1000
                    state["user_requirements"]["budget"] = budget
                    break
                except ValueError:
                    pass
        
        # Look for family size
        family_patterns = [
            r'family of (\d+)',
            r'(\d+) people',
            r'(\d+) of us',
        ]
        
        for pattern in family_patterns:
            matches = re.findall(pattern, message.lower())
            if matches:
                try:
                    state["user_requirements"]["family_size"] = int(matches[0])
                    break
                except ValueError:
                    pass
    
    def _build_context(self, state: ConversationState) -> str:
        return GeneralDesignAgent._build_context(self, state)


class FloorplanAgent(BaseAgent):
    """Handles floorplan planning and room layout"""
    
    def process(self, state: ConversationState) -> ConversationState:
        user_message = state["messages"][-1]["content"]
        context = self._build_context(state)
        floorplan_context = self._build_floorplan_context(state)
        
        system_prompt = f"""You are a residential floorplan specialist.
        
        User context:
        {context}
        
        Current floorplan requirements:
        {floorplan_context}
        
        IMPORTANT: Be proactive! Show initial designs with minimal information rather than asking many questions.
        
        Your approach:
        1. If user mentions ANY room needs (bedroom, bathroom, etc.) or house size - immediately offer to create an initial design
        2. Use reasonable defaults for missing information (assume 2000 sq ft if not specified, 2 floors for larger homes, etc.)
        3. Generate initial designs first, then ask for specific refinements
        4. Say something like: "I can create an initial floorplan design for you right now with these requirements, then we can refine it together"
        
        AVOID asking too many questions upfront. Users prefer to see a design first, then iterate. Ask user if they want more specifications after the design is shown."""
        
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=user_message)
        ]
        
        response = self.model.invoke(messages)
        
        # Extract floorplan information
        self._extract_floorplan_info(user_message, state)
        
        # Check if we have enough info for floorplan generation
        if self._check_floorplan_readiness(state):
            state["floorplan_ready"] = True
        
        state["messages"].append({
            "role": "assistant", 
            "content": response.content,
            "agent": "floorplan"
        })
        
        return state
    
    def _extract_floorplan_info(self, message: str, state: ConversationState):
        """Extract floorplan details from user message"""
        import re
        
        # Extract square footage
        sqft_patterns = [
            r'(\d{1,4})\s*(?:square feet|sq\.?\s*ft\.?|sqft)',
            r'(\d{1,4})\s*sq',
        ]
        
        for pattern in sqft_patterns:
            matches = re.findall(pattern, message.lower())
            if matches:
                try:
                    state["floorplan_requirements"]["total_sqft"] = int(matches[0])
                    break
                except ValueError:
                    pass
        
        # Extract number of floors
        floor_patterns = [
            r'(\d+)\s*(?:floor|story|stories)',
            r'(?:single|one)\s*(?:floor|story)' 
        ]
        
        for pattern in floor_patterns:
            if 'single' in message.lower() or 'one' in message.lower():
                state["floorplan_requirements"]["num_floors"] = 1
                break
            matches = re.findall(pattern, message.lower())
            if matches:
                try:
                    state["floorplan_requirements"]["num_floors"] = int(matches[0])
                    break
                except ValueError:
                    pass
        
        # Extract room information
        room_patterns = [
            (r'(\d+)\s*bedroom', 'bedroom'),
            (r'(\d+)\s*bathroom', 'bathroom'),
            (r'(\d+)\s*kitchen', 'kitchen'),
            (r'living room', 'living_room'),
            (r'dining room', 'dining_room'),
            (r'office', 'office'),
            (r'garage', 'garage'),
        ]
        
        for pattern, room_type in room_patterns:
            if room_type in ['living_room', 'dining_room', 'office', 'garage']:
                if room_type.replace('_', ' ') in message.lower():
                    # Check if this room type already exists
                    existing = [r for r in state["floorplan_requirements"]["rooms"] if r["type"] == room_type]
                    if not existing:
                        state["floorplan_requirements"]["rooms"].append({
                            "type": room_type,
                            "count": 1,
                            "min_size": None
                        })
            else:
                matches = re.findall(pattern, message.lower())
                if matches:
                    try:
                        count = int(matches[0])
                        # Check if this room type already exists
                        existing = [r for r in state["floorplan_requirements"]["rooms"] if r["type"] == room_type]
                        if existing:
                            existing[0]["count"] = count
                        else:
                            state["floorplan_requirements"]["rooms"].append({
                                "type": room_type,
                                "count": count,
                                "min_size": None
                            })
                    except ValueError:
                        pass
    
    def _check_floorplan_readiness(self, state: ConversationState) -> bool:
        """Check if we have enough info to generate a floorplan - now much more permissive"""
        reqs = state["floorplan_requirements"]
        user_reqs = state["user_requirements"]
        
        # Much more permissive - generate design with minimal information
        has_basic_rooms = len(reqs["rooms"]) >= 1  # Just one room type is enough
        has_any_size_info = (reqs["total_sqft"] is not None or 
                            reqs["lot_dimensions"] is not None or
                            user_reqs["budget"] is not None)
        has_basic_request = has_basic_rooms or has_any_size_info
        
        # Also trigger if user has mentioned any architectural terms
        recent_messages = state["messages"][-3:] if len(state["messages"]) >= 3 else state["messages"]
        architectural_keywords = ['bedroom', 'bathroom', 'kitchen', 'house', 'home', 'floor', 'room', 
                                'design', 'floorplan', 'layout', 'sqft', 'square feet', 'build']
        has_architectural_mention = any(keyword in msg.get("content", "").lower() 
                                      for msg in recent_messages 
                                      for keyword in architectural_keywords)
        
        # Generate floorplan much more readily - show initial design then iterate
        return has_basic_request or has_architectural_mention
    
    def _build_context(self, state: ConversationState) -> str:
        return GeneralDesignAgent._build_context(self, state)
    
    def _build_floorplan_context(self, state: ConversationState) -> str:
        reqs = state["floorplan_requirements"]
        context_parts = []
        
        if reqs["num_floors"]:
            context_parts.append(f"Floors: {reqs['num_floors']}")
        
        if reqs["total_sqft"]:
            context_parts.append(f"Total sq ft: {reqs['total_sqft']}")
        
        if reqs["lot_shape"]:
            context_parts.append(f"Lot shape: {reqs['lot_shape']}")
        
        if reqs["lot_dimensions"]:
            context_parts.append(f"Lot dimensions: {reqs['lot_dimensions']}")
        
        if reqs["rooms"]:
            rooms_str = ", ".join([f"{r['count']}x {r['type']}" for r in reqs["rooms"]])
            context_parts.append(f"Rooms: {rooms_str}")
        
        return "\n".join(context_parts) if context_parts else "No floorplan requirements collected yet."


class FloorplanGeneratorAgent(BaseAgent):
    """Generates thoughtfully designed architectural floorplans using design principles"""
    
    def process(self, state: ConversationState) -> ConversationState:
        # First, let the AI architect think through the design step-by-step
        design_analysis = self._analyze_design_requirements(state)
        
        # Create detailed floorplan simulation first
        detailed_simulation = self._create_detailed_simulation(state, design_analysis)
        
        # Store detailed floorplan information in shared state for other agents
        state["detailed_floorplan"]["design_analysis"] = design_analysis
        state["detailed_floorplan"]["detailed_rooms"] = detailed_simulation["rooms"]
        state["detailed_floorplan"]["structural_elements"] = detailed_simulation["structural"]
        state["detailed_floorplan"]["circulation_plan"] = detailed_simulation["circulation"]
        state["detailed_floorplan"]["lot_utilization"] = detailed_simulation["lot_usage"]
        state["detailed_floorplan"]["architectural_features"] = detailed_simulation["features"]
        
        # Store in agent memory for cross-agent access
        state["agent_memory"]["architectural_design"] = {
            "design_analysis": design_analysis,
            "detailed_simulation": detailed_simulation
        }
        
        # Generate detailed text-based floorplan with precise specifications
        # Removed SVG generation due to room overlap issues - focusing on precise text specifications
        simulation_response = self._format_simulation_response(detailed_simulation, design_analysis, state)
        
        state["messages"].append({
            "role": "assistant", 
            "content": simulation_response,
            "agent": "floorplan_generator"
        })
        
        state["floorplan_ready"] = False  # Reset for potential future generations
        
        return state
    
    def _analyze_design_requirements(self, state: ConversationState) -> Dict[str, Any]:
        """Use AI to analyze requirements and create architectural design strategy"""
        
        context = GeneralDesignAgent._build_context(self, state)
        floorplan_context = FloorplanAgent._build_floorplan_context(self, state)
        
        system_prompt = """You are a senior residential architect with 20+ years of experience. 
        Analyze the client's requirements and create a thoughtful architectural design strategy.
        
        Apply these key design principles:
        1. FUNCTIONAL ZONING: Separate public (living, dining, kitchen) from private (bedrooms) spaces
        2. CIRCULATION: Create efficient hallways and traffic flow, avoid dead ends
        3. NATURAL LIGHT: Orient main living spaces toward best light, minimize north-facing bedrooms
        4. PRIVACY: Bedrooms away from main entries, master suite separated from other bedrooms
        5. NOISE CONTROL: Kitchen/dining away from quiet bedrooms, buffer noisy areas
        6. VIEWS & OUTDOOR ACCESS: Living areas connect to outdoor spaces
        7. STRUCTURAL EFFICIENCY: Minimize long spans, stack plumbing, logical load paths
        8. BUILDING CODES: Ensure proper egress, room sizes, accessibility
        
        Think step-by-step through the design process:
        - Site orientation and entry location
        - Functional zoning strategy  
        - Circulation and hallway placement
        - Room adjacencies and relationships
        - Natural light optimization
        - Structural and mechanical considerations
        
        Return your analysis as a structured JSON with specific design decisions."""
        
        user_message = f"""
        CLIENT REQUIREMENTS:
        {context}
        
        FLOORPLAN REQUIREMENTS:
        {floorplan_context}
        
        Please analyze these requirements and provide a comprehensive architectural design strategy.
        Focus on creating a livable, functional home that follows good design principles.
        
        Respond with JSON format:
        {{
            "site_strategy": "orientation and entry approach",
            "zoning_concept": "how to separate public/private areas",
            "circulation_plan": "hallway and traffic flow strategy", 
            "room_relationships": "which rooms should be adjacent",
            "light_strategy": "natural light optimization approach",
            "design_priorities": ["priority1", "priority2", "priority3"],
            "floor_distribution": "how to distribute rooms across floors",
            "key_design_moves": ["move1", "move2", "move3"]
        }}
        """
        
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=user_message)
        ]
        
        response = self.model.invoke(messages)
        
        try:
            import json
            design_analysis = json.loads(response.content)
        except:
            # Fallback if JSON parsing fails
            design_analysis = {
                "site_strategy": "Standard residential approach with front entry",
                "zoning_concept": "Public spaces on ground floor, private spaces upstairs", 
                "circulation_plan": "Central hallway with efficient traffic flow",
                "room_relationships": "Kitchen adjacent to dining, living room central",
                "light_strategy": "South-facing living areas, east-facing bedrooms",
                "design_priorities": ["Functional layout", "Natural light", "Privacy"],
                "floor_distribution": "Common areas downstairs, bedrooms upstairs",
                "key_design_moves": ["Open concept main floor", "Private bedroom wing", "Efficient circulation"]
            }
        
        return design_analysis
    
    def _create_detailed_simulation(self, state: ConversationState, design_analysis: Dict[str, Any]) -> Dict[str, Any]:
        """Create detailed floorplan simulation with specific room dimensions and features"""
        
        reqs = state["floorplan_requirements"]
        user_reqs = state["user_requirements"]
        
        # Get project parameters with intelligent defaults based on available information
        budget = user_reqs.get('budget', 0) or 0
        family_size = user_reqs.get('family_size', 0) or 0
        
        # Estimate total square footage intelligently
        total_sqft = reqs.get("total_sqft")
        if not total_sqft:
            if budget > 800000:
                total_sqft = 2500  # Larger home for higher budget
            elif budget > 500000:
                total_sqft = 2000  # Standard home
            elif budget > 300000:
                total_sqft = 1500  # Smaller home
            elif (family_size or 0) >= 4:
                total_sqft = 2200  # Larger family needs more space
            elif (family_size or 0) >= 2:
                total_sqft = 1800  # Standard family
            else:
                total_sqft = 1600  # Default for unknown requirements
        
        # Estimate number of floors intelligently
        num_floors = reqs.get("num_floors")
        if not num_floors:
            if total_sqft > 2200:
                num_floors = 2  # Larger homes typically 2 floors
            else:
                num_floors = 1  # Default to single floor
        
        # Default lot dimensions
        lot_dims = reqs.get("lot_dimensions", "50x120 feet")
        
        # Generate basic room requirements if none provided
        rooms_list = reqs.get("rooms", [])
        if not rooms_list:
            # Create basic room requirements based on family size and square footage
            if (family_size or 0) >= 4 or total_sqft > 2000:
                rooms_list = [
                    {"type": "bedroom", "count": 3, "min_size": None},
                    {"type": "bathroom", "count": 2, "min_size": None},
                    {"type": "living_room", "count": 1, "min_size": None},
                    {"type": "kitchen", "count": 1, "min_size": None},
                    {"type": "dining_room", "count": 1, "min_size": None}
                ]
            else:
                rooms_list = [
                    {"type": "bedroom", "count": 2, "min_size": None},
                    {"type": "bathroom", "count": 1, "min_size": None},
                    {"type": "living_room", "count": 1, "min_size": None},
                    {"type": "kitchen", "count": 1, "min_size": None}
                ]
        
        # Simulate detailed architectural design process with safe formatting
        budget_str = f"${budget:,.0f}" if budget else "Estimated based on size"
        family_str = f"{family_size} people" if family_size else "Estimated based on requirements"
        
        context = f"""
        Project: {total_sqft} sq ft, {num_floors} floor home (using intelligent defaults where needed)
        Family: {family_str}
        Budget: {budget_str} CAD
        Lot: {lot_dims}
        Rooms needed: {', '.join([f"{r['count']}x {r['type']}" for r in rooms_list])}
        
        Design Strategy: {design_analysis.get('zoning_concept', 'Functional zoning')}
        Note: Generated with available information - user can refine any aspect
        """
        
        system_prompt = """You are a senior architect creating a detailed floorplan simulation.
        
        IMPORTANT: Create excellent designs even with limited information. Use professional judgment 
        to fill in missing details with reasonable, well-designed solutions.
        
        Based on the project requirements (some may be estimated), create specific room dimensions, placements, and features.
        Consider:
        - Optimal room sizes for function and comfort
        - Traffic flow and circulation
        - Natural light and orientation  
        - Structural efficiency
        - Building codes and accessibility
        - Cost-effective design
        - Standard residential design best practices
        
        If information is missing or estimated, use your architectural expertise to create
        a functional, livable design that the user can then refine.
        
        Provide detailed specifications for each room including exact dimensions, placement rationale,
        special features, and architectural elements.
        
        Return detailed JSON simulation data."""
        
        user_message = f"""
        {context}
        
        Create a detailed architectural simulation with:
        
        1. Specific room dimensions and square footage
        2. Room placement rationale (why each room is positioned where it is)
        3. Special features and built-ins for each room
        4. Structural elements (load-bearing walls, beams, etc.)
        5. Circulation plan (hallways, stairs, traffic flow)
        6. Lot utilization (building placement, setbacks, outdoor spaces)
        7. Architectural features (windows, doors, ceiling heights, etc.)
        
        Return JSON:
        {{
            "rooms": [
                {{
                    "type": "room_type",
                    "label": "Room Name", 
                    "width": width_in_feet,
                    "height": height_in_feet,
                    "sqft": square_footage,
                    "floor": floor_number,
                    "placement_rationale": "why this room is here",
                    "features": ["feature1", "feature2"],
                    "natural_light": "light_description",
                    "finishes": "finish_specifications"
                }}
            ],
            "structural": [
                {{
                    "type": "structural_element",
                    "description": "element description",
                    "location": "where it is",
                    "purpose": "structural purpose"
                }}
            ],
            "circulation": {{
                "hallways": [{{
                    "location": "hallway location",
                    "width": width_feet,
                    "purpose": "circulation purpose"
                }}],
                "stairs": {{
                    "type": "stair_type",
                    "location": "stair_location", 
                    "width": width_feet
                }},
                "traffic_flow": "description of movement patterns"
            }},
            "lot_usage": {{
                "building_footprint": "building size and placement",
                "setbacks": "distance from property lines",
                "outdoor_spaces": ["outdoor space descriptions"],
                "parking": "parking arrangement",
                "landscaping": "landscape considerations"
            }},
            "features": ["special architectural features"]
        }}
        """
        
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=user_message)
        ]
        
        response = self.model.invoke(messages)
        
        try:
            import json
            simulation_data = json.loads(response.content)
            return simulation_data
        except:
            # Fallback simulation if JSON parsing fails
            return {
                "rooms": [
                    {
                        "type": "living_room",
                        "label": "Living Room",
                        "width": 18,
                        "height": 16,
                        "sqft": 288,
                        "floor": 1,
                        "placement_rationale": "Central location for family gathering",
                        "features": ["Fireplace", "Large windows", "Open to kitchen"],
                        "natural_light": "South-facing windows for optimal light",
                        "finishes": "Hardwood flooring, painted walls"
                    }
                ],
                "structural": [
                    {
                        "type": "load_bearing_wall",
                        "description": "Central beam supporting second floor",
                        "location": "Between living and dining areas",
                        "purpose": "Structural support for upper floor"
                    }
                ],
                "circulation": {
                    "hallways": [
                        {
                            "location": "Central hallway connecting main areas",
                            "width": 4,
                            "purpose": "Primary circulation between public spaces"
                        }
                    ],
                    "stairs": {
                        "type": "straight_stair",
                        "location": "Near entrance",
                        "width": 3.5
                    },
                    "traffic_flow": "Efficient movement between all areas"
                },
                "lot_usage": {
                    "building_footprint": f"{(total_sqft or 2000) // (num_floors or 1)} sq ft footprint",
                    "setbacks": "Standard residential setbacks",
                    "outdoor_spaces": ["Front yard", "Backyard patio"],
                    "parking": "Driveway and garage",
                    "landscaping": "Lawn and foundation plantings"
                },
                "features": ["Open concept design", "High ceilings", "Large windows"]
            }
    
    def _format_simulation_response(self, simulation: Dict[str, Any], design_analysis: Dict[str, Any], state: ConversationState) -> str:
        """Format the detailed simulation into a comprehensive response with precise design parameters"""
        
        reqs = state["floorplan_requirements"]
        user_reqs = state["user_requirements"]
        
        # Safe formatting for numeric values
        total_sqft = reqs.get('total_sqft', 0) or 0
        budget = user_reqs.get('budget', 0) or 0
        lot_dimensions = reqs.get('lot_dimensions') or '50x120 feet'
        
        response = f"""🏗️ **DETAILED FLOORPLAN WITH EXACT SPECIFICATIONS** 🏗️

**📋 PROJECT PARAMETERS:**
• Total Area: {total_sqft:,} sq ft
• Building Footprint: {(total_sqft or 2000) // (reqs.get('num_floors', 1) or 1):,} sq ft per floor
• Floors: {reqs.get('num_floors', 1)}
• Lot Size: {lot_dimensions}
• Family Size: {user_reqs.get('family_size') or 'TBD'}
• Budget: ${budget:,.0f} CAD

**🎯 DESIGN APPROACH:**
• Primary Strategy: {design_analysis.get('zoning_concept', 'Functional zoning with clear separation of public and private areas')}
• Traffic Flow: {design_analysis.get('circulation_plan', 'Central circulation spine minimizing hallway space')}
• Natural Light: {design_analysis.get('light_strategy', 'South-facing living areas, east-facing bedrooms')}

📐 **PRECISE ROOM LAYOUT - FLOOR BY FLOOR:**
"""
        
        # Group rooms by floor and add precise specifications
        rooms = simulation.get("rooms", [])
        floors = {}
        total_simulated_sqft = 0
        
        for room in rooms:
            floor_num = room.get("floor", 1)
            if floor_num not in floors:
                floors[floor_num] = []
            floors[floor_num].append(room)
            
            width = room.get("width", 0) or 0
            height = room.get("height", 0) or 0
            sqft = room.get("sqft", width * height) or 0
            total_simulated_sqft += sqft
        
        # Display each floor with precise measurements
        for floor_num in sorted(floors.keys()):
            floor_rooms = floors[floor_num]
            floor_sqft = sum((r.get("sqft") or ((r.get("width") or 0) * (r.get("height") or 0))) for r in floor_rooms)
            
            response += f"\n🏠 **FLOOR {floor_num} LAYOUT** ({floor_sqft:,} sq ft)\n"
            response += "=" * 50 + "\n"
            
            for room in floor_rooms:
                width = room.get("width", 0) or 0
                height = room.get("height", 0) or 0
                sqft = room.get("sqft") or (width * height)
                
                response += f"""
🏠 **{room.get('label', 'Room').upper()}**
   📏 Exact Dimensions: {width}' × {height}' = {sqft} sq ft
   📍 Location: {room.get('placement_rationale', 'Optimal positioning')}
   🚪 Access: {room.get('access', 'Standard door access')}
   🪟 Windows: {room.get('windows', 'Natural light from south/east')}
   🔌 Electrical: {room.get('electrical', '4-6 outlets, overhead lighting')}
   🏗️ Ceiling: {room.get('ceiling_height', '9')}' height
   🎨 Finishes: {room.get('finishes', 'Hardwood floors, painted walls')}"""
                
                # Add room-specific details
                room_type = room.get('type', '').lower()
                if 'kitchen' in room_type:
                    response += f"\n   🍳 Kitchen Features: Island {room.get('island_size', '4x8')}', cabinets along {room.get('cabinet_walls', '2 walls')}"
                elif 'bathroom' in room_type:
                    response += f"\n   🚿 Bathroom Layout: {room.get('bathroom_layout', '3-piece with shower/tub combo')}"
                elif 'bedroom' in room_type:
                    response += f"\n   🛏️ Bedroom Setup: {room.get('bed_size', 'Queen')} bed, {room.get('closet', 'walk-in closet')}"
                
                response += "\n"
        
        response += f"\n📊 **TOTAL LIVING SPACE: {total_simulated_sqft:,} sq ft**\n"
        
        # Add precise structural and dimensional information
        circulation = simulation.get("circulation", {})
        if circulation:
            response += "\n🚶 **CIRCULATION SPECIFICATIONS:**\n"
            
            hallways = circulation.get("hallways", [])
            for hall in hallways:
                response += f"• {hall.get('location', 'Main hallway')}: {hall.get('width', 4)}' wide × {hall.get('length', 20)}' long\n"
            
            stairs = circulation.get("stairs", {})
            if stairs and (reqs.get('num_floors') or 1) > 1:
                response += f"• Staircase: {stairs.get('type', 'Straight run')} - {stairs.get('width', 3.5)}' wide, {stairs.get('rise', '7')}\" rise, {stairs.get('run', '11')}\" run\n"
        
        # Add precise lot utilization
        lot_usage = simulation.get("lot_usage", {})
        if lot_usage:
            response += "\n🏡 **LOT UTILIZATION PLAN:**\n"
            try:
                # Parse lot dimensions safely
                if 'x' in lot_dimensions:
                    lot_parts = lot_dimensions.lower().replace('feet', '').replace('ft', '').strip().split('x')
                    lot_width = int(lot_parts[0].strip())
                    lot_depth = int(lot_parts[1].strip())
                else:
                    lot_width, lot_depth = 50, 120  # Default
                
                building_width = lot_width - 10  # 5' setbacks each side
                footprint_sqft = (total_sqft or 2000) // (reqs.get('num_floors', 1) or 1)
                building_depth = min(footprint_sqft // max(building_width, 1), lot_depth - 35)  # Ensure it fits
                
                response += f"• Building Position: Centered on lot with 5' side setbacks\n"
                response += f"• Building Footprint: {building_width}' × {building_depth}' ({building_width * building_depth:,} sq ft)\n"
                response += f"• Front Setback: 20' from street\n"
                response += f"• Rear Setback: 15' from back property line\n"
                response += f"• Driveway: 12' wide × 20' deep\n"
                response += f"• Front Yard: {lot_width}' × 20'\n"
                response += f"• Back Yard: {lot_width}' × {lot_depth - building_depth - 35}'\n"
            except (ValueError, IndexError):
                # Fallback if lot dimensions parsing fails
                response += f"• Building Position: Optimal placement on {lot_dimensions} lot\n"
                response += f"• Standard setbacks and yard arrangements\n"
        
        # Add precise construction specifications
        structural = simulation.get("structural", [])
        if structural:
            response += "\n🏗️ **STRUCTURAL SPECIFICATIONS:**\n"
            for element in structural:
                response += f"• {element.get('type', 'Structural element')}: {element.get('specifications', element.get('description', 'Standard construction'))}\n"
        
        response += f"""

🔧 **CONSTRUCTION SPECIFICATIONS:**
• Foundation: {simulation.get('foundation_type', 'Concrete slab-on-grade or full basement')}
• Framing: {simulation.get('framing', '2x6 wood frame construction, 16" O.C.')}
• Insulation: {simulation.get('insulation', 'R-24 walls, R-50 attic, R-12 basement')}
• Electrical: {simulation.get('electrical_specs', '200-amp panel, GFCI in wet areas')}
• Plumbing: {simulation.get('plumbing_specs', 'PEX supply lines, ABS drain lines')}
• HVAC: {simulation.get('hvac', 'Forced air gas furnace, central air conditioning')}

📐 **DETAILED DRAWING GUIDE:**

**STEP-BY-STEP DRAWING INSTRUCTIONS:**
1. **Scale Setup**: Use 1/4" = 1' scale (1:48) for standard architectural drawings
2. **Grid Method**: Start with a grid - each square = 2 feet for easy measurement
3. **Exterior Walls**: Draw building perimeter first using lot placement dimensions above
4. **Room Layout**: Use exact room dimensions provided - each room precisely measured
5. **Wall Thickness**: Standard interior walls = 4", exterior walls = 6"
6. **Door Openings**: Standard doors = 3' wide, mark swing direction
7. **Window Placement**: Based on natural light specifications for each room

**ARCHITECTURAL STANDARDS REFERENCE:**
• Minimum Room Sizes: Bedrooms 70+ sq ft, Bathrooms 40+ sq ft, Kitchens 70+ sq ft
• Hallway Widths: Main hallways 4' minimum, secondary 3' minimum  
• Ceiling Heights: Standard 9', can vary by room as specified
• Stair Requirements: 7" max rise, 11" min run, 36" min width

**CAD/DRAFTING TIPS:**
• All dimensions in feet and inches (e.g., 14'-6")
• Room labels show exact square footage for verification
• Use dimensions provided to create proportionally accurate drawings
• Each room placement includes architectural reasoning for optimal layout

🔄 **DESIGN REFINEMENT:**
This is your INITIAL design based on the information provided. You can now:
• Request changes to any room sizes, locations, or features
• Add or remove rooms as needed
• Adjust the overall layout or style
• Modify lot placement or outdoor spaces
• Update any specifications or construction details

Just tell me what you'd like to change and I'll create an updated design!

**NEXT STEPS:**
• Review this initial design and request any modifications
• Budget estimation based on exact square footages and construction details
• Structural engineering review using provided foundation and framing specs
• Building permit applications with detailed room schedules and calculations"""
        
        return response


class RegulationAgent(BaseAgent):
    """Provides Montreal building code guidance and permit requirements with official source citations"""
    
    def process(self, state: ConversationState) -> ConversationState:
        user_message = state["messages"][-1]["content"]
        
        # Build safe context from existing data
        context = self._build_regulation_context(state)
        
        system_prompt = f"""You are a Montreal building code specialist and permit consultant with expertise in Quebec construction regulations.

        Provide detailed guidance on:
        - Montreal/Quebec building codes and zoning requirements
        - Montreal/Quebec construction code
        - Properly show the user the code and original source of the information
        - Required permits for residential construction and renovations
        - Setback requirements and lot coverage rules
        - Room size minimums and ceiling height requirements
        - Egress and safety code requirements
        - Electrical, plumbing, and HVAC code compliance
        - Timeline for permit applications and approvals
        - Required professional services (architect/engineer stamps)
        - Inspection schedules and requirements
        - Neighbor notification and variance procedures

        Current project context:
        {context}

        IMPORTANT: 
        - Provide your knowledge of Montreal building regulations and Quebec Construction Code
        - Include specific regulation names, bylaw numbers, and code sections when you know them
        - Mention key requirements like minimum room sizes, setback distances, permit types
        - DO NOT provide website URLs or links - focus on the actual regulatory content
        - Be specific about Montreal zoning requirements, permit processes, and timelines
        - Include information about professional requirements (architect stamps, etc.)

        Format your response with clear regulatory information that users can reference when contacting Montreal building officials."""
        
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=user_message)
        ]
        
        response = self.model.invoke(messages)
        
        # Store regulation guidance in shared memory
        state["agent_memory"]["regulation_guidance"] = {
            "consultation": response.content,
            "topics_covered": self._extract_regulation_topics(user_message)
        }
        
        formatted_response = f"""📋 **MONTREAL BUILDING REGULATIONS & PERMITS** 📋

{response.content}

**🔍 NEXT STEPS:**
• Contact City of Montreal building department to verify current requirements
• Visit montreal.ca and search for "building permits" and "zoning" for official information
• Consult with licensed architect/engineer for official compliance review
• Obtain current application forms from Montreal building department

**⚖️ DISCLAIMER:** This guidance is educational and based on general knowledge of Montreal building regulations. Official compliance must be verified with current City of Montreal building department requirements and Quebec Construction Code. Always consult licensed professionals for your specific project."""

        state["messages"].append({
            "role": "assistant",
            "content": formatted_response,
            "agent": "regulation"
        })
        
        return state
    
    def _build_regulation_context(self, state: ConversationState) -> str:
        """Build context safely without risky formatting"""
        context_parts = []
        
        # Safe access to floorplan requirements
        floor_reqs = state["floorplan_requirements"]
        if floor_reqs.get("total_sqft"):
            context_parts.append(f"House Size: {floor_reqs['total_sqft']} sq ft")
        
        if floor_reqs.get("num_floors"):
            context_parts.append(f"Floors: {floor_reqs['num_floors']} stories")
            
        if floor_reqs.get("lot_dimensions"):
            context_parts.append(f"Lot: {floor_reqs['lot_dimensions']}")
        
        # Safe access to room information
        if floor_reqs.get("rooms"):
            room_summary = []
            for room in floor_reqs["rooms"]:
                room_type = room.get("type", "room")
                room_count = room.get("count", 1)
                room_summary.append(f"{room_count} {room_type}")
            if room_summary:
                context_parts.append(f"Rooms: {', '.join(room_summary)}")
        
        # Safe access to user location
        user_reqs = state["user_requirements"]
        location = user_reqs.get("location", "Montreal")
        context_parts.append(f"Location: {location}")
        
        # Check if detailed floorplan exists
        detailed_fp = state["detailed_floorplan"]
        if detailed_fp.get("detailed_rooms"):
            context_parts.append("Status: Detailed floorplan completed")
        
        # Check project type
        budget = user_reqs.get("budget")
        if budget and budget > 0:
            context_parts.append("Project: New construction with established budget")
        else:
            context_parts.append("Project: Planning phase")
        
        if context_parts:
            return "\n".join([f"• {part}" for part in context_parts])
        else:
            return "• New residential project in Montreal area"
    
    def _extract_regulation_topics(self, message: str) -> list:
        """Extract regulation topics from user message"""
        topics = []
        message_lower = message.lower()
        
        topic_keywords = {
            "permits": ["permit", "approval", "application"],
            "zoning": ["zoning", "setback", "lot coverage"],
            "building_code": ["building code", "code", "regulation", "requirement"],
            "electrical": ["electrical", "wiring", "panel"],
            "plumbing": ["plumbing", "water", "sewer"],
            "structural": ["structural", "foundation", "beam"],
            "safety": ["safety", "egress", "fire", "emergency"],
            "inspections": ["inspection", "inspector", "review"]
        }
        
        for topic, keywords in topic_keywords.items():
            if any(keyword in message_lower for keyword in keywords):
                topics.append(topic)
        
        return topics if topics else ["general_inquiry"]