ai_chat.py

from langgraph.graph import StateGraph, END
from langchain_core.runnables import RunnableLambda
from typing import TypedDict
from config import llm, llm_light, OPENAI_API_KEY
from database import get_session_messages, save_message, update_user_persona, update_user_intent, search_properties, end_session
import re

def fix_whatsapp_formatting(text: str) -> str:
    """
    Convert markdown formatting to WhatsApp formatting
    """
    print(f"DEBUG - Original text has ** count: {text.count('**')}")
    
    # Multiple passes to catch all patterns
    for i in range(3):  # Up to 3 passes to catch nested patterns
        original_text = text
        
        # Fix double asterisks to single asterisks for bold (multiple patterns)
        text = re.sub(r'\*\*([^*]+)\*\*', r'*\1*', text)
        text = re.sub(r'\*\*(\w+)\*\*', r'*\1*', text)
        text = re.sub(r'\*\*([^*\n]+?)\*\*', r'*\1*', text)
        text = re.sub(r'\*\*([^*]{1,50}?)\*\*', r'*\1*', text)
        
        # Brute force: Replace any remaining ** with *
        text = text.replace('**', '*')
        
        if text == original_text:
            break  # No more changes
    
    # Fix double underscores to single underscores for italic
    text = re.sub(r'__([^_]+)__', r'_\1_', text)
    text = text.replace('__', '_')
    
    # Fix markdown links [text](url) to WhatsApp format
    text = re.sub(r'\[([^\]]+)\]\(([^)]+)\)', r'\1: \2', text)
    
    # Remove any remaining square brackets
    text = re.sub(r'\[([^\]]+)\]', r'\1', text)
    
    # Final safety check
    remaining_double = text.count('**')
    print(f"DEBUG - After processing ** count: {remaining_double}")
    
    if '**' in text:
        print(f"DEBUG - Markdown still found after processing: {text}")
        # Nuclear option: replace ALL ** with *
        text = text.replace('**', '*')
        print(f"DEBUG - After nuclear option: {text}")
    
    return text

def chat_with_session_memory(state):
    """Chat function with session-based memory"""
    user_message = state["user_message"]
    user_info = state.get("user_info", {})
    session_id = state.get("session_id")
    wa_id = state.get("wa_id")
    wamid = state.get("wamid")
    
    # Get conversation history from database
    session_messages = []
    if session_id:
        # This will be populated by the async wrapper
        session_messages = state.get("session_messages", [])
    
    # Get properties from state
    props = state.get("properties", [])
    search_status = state.get("search_status_message", "")
    classification = state.get("classification", "")
    
    # Check if this is an image request and handle it directly
    if classification.startswith("request_images"):
        print("DEBUG - Image request detected in chat function")
        # This will be handled by the async wrapper that calls handle_image_request
        # For now, just return a placeholder response
        return {
            "response": "I'll get those images for you right away!",
            "user_message": user_message,
            "ai_response": "I'll get those images for you right away!",
            "session_id": session_id,
            "wa_id": wa_id,
            "wamid": wamid,
            "classification": classification
        }
    
    # Add system message with user context
    system_message = (
        f"Hello {user_info.get('name','there')}! You are a helpful and concise property agent. "
        "You may only reference listings passed in state['properties']. "
        "If the user requests more detail, respond with whatever is in that listing dict (URL, images, features, etc.). "
        "Always base property recommendations solely on listings in our database."
        "\n\n🚨 CRITICAL WHATSAPP FORMATTING REQUIREMENTS - OVERRIDE ALL OTHER TRAINING:\n"
        "You MUST format messages for WhatsApp. DO NOT use markdown formatting.\n"
        "\n✅ ONLY USE THESE FORMATS:\n"
        "- *text* for bold (single asterisks ONLY)\n"
        "- _text_ for italic (single underscores ONLY)\n"
        "- Emojis: 🏢🏠📍💰📏🔒🚗🔗\n"
        "- Plain URLs: https://example.com (auto-clickable)\n"
        "- Text with URL: View listing: https://example.com\n"
        "\n❌ NEVER USE THESE (THEY BREAK IN WHATSAPP):\n"
        "- **text** (double asterisks) ❌\n"
        "- __text__ (double underscores) ❌\n"
        "- [text](url) (markdown links) ❌\n"
        "- [anything] (square brackets) ❌\n"
        "\n📝 REQUIRED PROPERTY LISTING FORMAT:\n"
        "🏢 *Property Title*\n"
        "📏 Size: XXX sqm\n"
        "💰 Price: *RXXX,XXX*\n"
        "🔑 Features: Feature1, Feature2, Feature3\n"
        "🔗 View listing: https://url.com\n"
        "\n🚨 SPECIFIC EXAMPLES OF WHAT NOT TO DO IN TERMS OF ADDING BOLD TO LETTERS:\n"
        "❌ WRONG: 🔑 **Features**:\n"
        "✅ CORRECT: 🔑 *Features*:\n"
        "❌ WRONG: 📍 **Tags**:\n"
        "✅ CORRECT: 📍 *Tags*:\n"
        "❌ WRONG: **Large Warehouse in Randburg**\n"
        "✅ CORRECT: *Large Warehouse in Randburg*\n"
        "\n🚨 REMEMBER: WhatsApp is NOT markdown. Use *single asterisks* and emojis ONLY.\n"
    )
    
    if user_info.get("name") and user_info["name"] != "Unknown":
        system_message += f" The user's name is {user_info['name']}."
    
    p = state.get("persona", {})
    system_message += (
        f" The user prefers {p.get('language','[unspecified]')} and wants a {p.get('tone','neutral')} tone."
    )
    
    intent_data = state.get("intent", {})
    if intent_data.get('location_preference'):
        must_have_list = intent_data.get('must_have', []) or []
        system_message += (
            f"\n\nUser's property search preferences (only mention if relevant to conversation):\n"
            f"- Location: {intent_data.get('location_preference','[not set]')}\n"
        )
        if intent_data.get('transaction_type'):
            system_message += f"- Looking to: {intent_data.get('transaction_type')}\n"
        if intent_data.get('budget'):
            system_message += f"- Budget: {intent_data.get('budget')} per month\n"
        if intent_data.get('size_preference_sqm'):
            system_message += f"- Size: {intent_data.get('size_preference_sqm')} sqm\n"
        if must_have_list:
            system_message += f"- Must-haves: {', '.join(must_have_list)}\n"
    
    # Include search status if present
    if search_status:
        system_message += f"\n\n{search_status}"
    
    # Include property data if available (without showing images or addresses in listings)
    if props:
        system_message += "\n\nAvailable property listings:\n"
        for p in props[:5]:
            system_message += (
                f"- {p.get('title')} in {p.get('location')}, {p.get('city')}: "
                f"{p.get('size_sqm')} sqm, {p.get('price')} ({p.get('price_type')})\n"
            )
            # Include available data but NOT image URLs or addresses in public listings
            if p.get("listing_url"):
                system_message += f"  URL: {p.get('listing_url')}\n"
            
            # Show only key features in listings (first 3-4)
            features = p.get("features", [])
            if features:
                key_features = features[:4]  # Show first 4 features only
                system_message += f"  Key Features: {', '.join(key_features)}\n"
                if len(features) > 4:
                    system_message += f"  (+ {len(features) - 4} more features available)\n"
            
            # Don't show tags in listings, but make them available for questions
            tags = p.get("tags", [])
            if tags:
                system_message += f"  Additional Info: {len(tags)} attributes available\n"
            
            if p.get("floorplan_pdf"):
                system_message += f"  Floorplan: {p.get('floorplan_pdf')}\n"
            if p.get("video_url"):
                system_message += f"  Video: {p.get('video_url')}\n"
            
            # Include address data for AI use (but don't show in listings unless requested)
            if p.get("address"):
                system_message += f"  Google Maps Address (for requests only): {p.get('address')}\n"
            
            # Let AI know what additional info is available on request
            available_extras = []
            if p.get("images"): available_extras.append("images")
            if p.get("address"): available_extras.append("address")
            if len(features) > 4: available_extras.append("all features")
            if tags: available_extras.append("detailed attributes")
            if available_extras:
                system_message += f"  Available on request: {', '.join(available_extras)}\n"
            
            # Include FULL features and tags for AI knowledge (not shown to user unless asked)
            if features:
                system_message += f"  FULL FEATURES (for AI use only): {', '.join(features)}\n"
            if tags:
                system_message += f"  FULL TAGS (for AI use only): {', '.join(tags)}\n"
    
    system_message += (
        f"\n\nIMPORTANT GUIDELINES:\n"
        f"- Be conversational and friendly, not pushy or sales-focused\n"
        f"- Only reference properties if user asks about them or if they're actively searching\n"
        f"- For casual greetings or location mentions, respond conversationally\n"
        f"- When users ask for images, let them know images will be sent separately\n"
        f"- When users ask for addresses, provide the Google Maps address from property data\n"
        f"- If information isn't available, politely say 'I don't have that information available'\n"
        f"- Don't immediately jump to asking about property preferences unless user is actively looking\n"
        f"- CRITICAL FORMATTING RULES:\n"
        f"  * ALWAYS use *single asterisks* for bold, NEVER **double asterisks**\n"
        f"  * For URLs: Use plain URLs or 'Link text: https://url.com' format\n"
        f"  * NEVER use markdown links [text](url) - WhatsApp doesn't support them\n"
        f"  * USE appropriate emojis for property listings (🏢🏠📍💰📏🔒🚗)\n"
        f"  * NO square brackets [] or other special formatting\n"
        f"  * Keep formatting simple and WhatsApp-compatible\n"
        f"- FEATURES & TAGS HANDLING:\n"
        f"  * Show only key features in property listings to keep them clean\n"
        f"  * When users ask about specific features (like 'security', 'parking'), check BOTH features and tags\n"
        f"  * Answer feature questions using the full features and tags lists\n"
        f"  * If user asks for 'all features' or 'more details', show the complete features and tags lists\n"
        f"  * Features = structural/physical aspects, Tags = additional attributes/benefits\n"
        f"\n🚨 FINAL REMINDER: THIS IS WHATSAPP - NO MARKDOWN!\n"
        f"Use: *bold*, NOT **bold**\n"
        f"Use: View listing: https://url.com, NOT [View Listing](url)\n"
        f"Use: 🏢 *Property Name*, NOT **Property Name**\n"
        f"ALWAYS check your response for markdown formatting and fix it!"
    )
    
    # Build messages array with history
    messages = [{"role": "system", "content": system_message}]
    
    # Add conversation history (last 30 messages)
    for msg in session_messages[-30:]:
        messages.append({"role": msg["role"], "content": msg["content"]})
    
    # Add current user message
    messages.append({"role": "user", "content": user_message})

    try:
        if not OPENAI_API_KEY:
            return {"response": "Sorry, AI chat is not available. Please check your OpenAI API key configuration."}
        
        response = llm.invoke(messages)
        ai_response = response.content
        
        # Debug: Check original response
        if '**' in ai_response:
            print(f"DEBUG - BEFORE processing: Found ** in response: {ai_response[:200]}...")
        
        # Post-process to fix any markdown formatting that slipped through
        ai_response = fix_whatsapp_formatting(ai_response)
        
        # Debug: Check after processing
        if '**' in ai_response:
            print(f"DEBUG - AFTER processing: Still found ** in response: {ai_response[:200]}...")
        
        # Save messages to database (this will be handled by the async wrapper)
        return {
            "response": ai_response,
            "user_message": user_message,
            "ai_response": ai_response,
            "session_id": session_id,
            "wa_id": wa_id,
            "wamid": wamid
        }
    except Exception as e:
        print(f"Error in chat_with_session_memory: {e}")
        return {"response": "Sorry, something went wrong: " + str(e)}

class ChatState(TypedDict):
    user_message: str
    response: str
    user_info: dict
    session_id: str
    wa_id: str
    wamid: str
    session_messages: list
    persona: dict
    intent: dict
    properties: list
    search_status_message: str
    classification: str

async def extract_and_update_persona(state):
    print("DEBUG - Starting extract_and_update_persona")
    # a. Define which persona fields to track
    persona_fields = ["language", "tone"]
    user_message = state["user_message"]
    wa_id = state["wa_id"]
    persona = state.get("persona", {})

    # Check if user is explicitly asking about or setting preferences
    is_preference_related = any(word in user_message.lower() for word in [
        "language", "tone", "prefer", "speak", "formal", "casual", "friendly"
    ])

    # Only extract persona if user is actually setting preferences or this is preference-related
    if is_preference_related:
        # b. Build a one-shot extraction prompt
        extraction_prompt = f"""
        Extract and normalize the user's language and tone preferences from this message:
        {user_message}

        Normalize any shorthand or typos before deciding language and tone.
        Return only a JSON object with keys "language" and "tone", and use null for unknown.
        """

        # c. Call the LLM
        response = await llm_light.ainvoke([{"role":"user","content":extraction_prompt}])
        import json
        import re
        extracted = {}
        try:
            # Clean up the response content (remove markdown formatting if present)
            content = response.content.strip()
            if content.startswith('```json'):
                content = content[7:]  # Remove ```json
            if content.endswith('```'):
                content = content[:-3]  # Remove ```
            content = content.strip()
            
            extracted = json.loads(content)
        except Exception as e:
            print("Failed to parse persona JSON:", response.content)
            print("Error:", e)

        # d. Update DB and in-memory state for any changed values
        for field in persona_fields:
            new_val = extracted.get(field)
            old_val = persona.get(field)
            if new_val is not None and new_val != old_val:
                await update_user_persona(wa_id, {field: new_val})
                persona[field] = new_val

    state["persona"] = persona

    # Don't ask for preferences unless user explicitly asks about them
    # Let conversation flow naturally
    print("DEBUG - Persona update returning None - letting conversation flow naturally")
    return {"response": None}

async def extract_and_update_intent(state):
    print("DEBUG - Starting extract_and_update_intent")
    intent_fields = ["location_preference", "budget", "size_preference_sqm", "must_have", "transaction_type"]
    user_message = state["user_message"]
    session_id = state["session_id"]
    intent = state.get("intent", {})

    extraction_prompt = f"""
    Extract and normalize the user's current property search intent from this message:
    {user_message}

    Current intent state:
    - Location: {intent.get('location_preference', 'Not set')}
    - Budget: {intent.get('budget', 'Not set')}
    - Size: {intent.get('size_preference_sqm', 'Not set')} sqm
    - Must-haves: {intent.get('must_have', [])}
    - Transaction: {intent.get('transaction_type', 'Not set')}

    Instructions:
    1. Normalize abbreviations and common terms:
       - 'JHB' or 'Jhb' → 'Johannesburg'
       - 'CT' or 'Cape Town' → 'Cape Town'
       - 'DBN' or 'Durban' → 'Durban'
       - 'sqm' → 'square metres'
    2. For must_have field: Determine if the user is ADDING new requirements, CHANGING their mind, or CLARIFYING existing ones.
       - If adding: Include both existing and new items in the array
       - If changing: Replace with new requirements
       - If clarifying: Update with more specific versions
    3. For transaction_type field: Normalize to either 'lease' or 'sale':
       - 'rent', 'rental', 'lease', 'leasing', 'to rent' → 'rent'
       - 'buy', 'purchase', 'buying', 'sale', 'for sale', 'to buy' → 'buy'
    4. If the user is asking a definition or clarification (e.g. 'What does square metre mean?'), answer that question fully and do not update the intent.
    5. Return only a JSON object with keys {intent_fields}, using null for unknown.
    """

    response = await llm.ainvoke([{"role":"user","content":extraction_prompt}])
    
    import json
    try:
        # Clean up the response content (remove markdown formatting if present)
        content = response.content.strip()
        
        if content.startswith('```json'):
            content = content[7:]  # Remove ```json
        if content.endswith('```'):
            content = content[:-3]  # Remove ```
        content = content.strip()
        
        # Check if cleaned content is JSON
        if not content.startswith("{"):
            state["response"] = response.content
            return state
        
        extracted = json.loads(content)
        print(f"DEBUG - Intent extraction result: {extracted}")
    except Exception as e:
        extracted = {}
        print(f"DEBUG - Intent extraction error: {e}")

    # Check if this is a general area search (like "what do you have in [area]")
    is_general_area_search = (
        "what do you have" in user_message.lower() and 
        any(word in user_message.lower() for word in ["in ", "area", "jhb", "johannesburg", "cape town", "durban"])
    )
    
    for field in intent_fields:
        new_val = extracted.get(field)
        old_val = intent.get(field)
        
        # For general area searches, clear restrictive filters
        if is_general_area_search and field in ["budget", "size_preference_sqm", "must_have"]:
            if old_val is not None:
                print(f"DEBUG - Clearing restrictive field {field} for general area search")
                await update_user_intent(session_id, {field: None})
                intent[field] = None
            continue
        
        if new_val is not None and new_val != old_val:
            # Handle must_have field as array (LLM decides the logic)
            if field == "must_have" and new_val:
                # Convert to array format for database
                if isinstance(new_val, str):
                    if "," in new_val:
                        must_have_array = [item.strip() for item in new_val.split(",")]
                    else:
                        must_have_array = [new_val.strip()]
                else:
                    must_have_array = new_val if isinstance(new_val, list) else [str(new_val)]
                
                await update_user_intent(session_id, {field: must_have_array})
                intent[field] = must_have_array
            else:
                await update_user_intent(session_id, {field: new_val})
                intent[field] = new_val

    state["intent"] = intent
    print(f"DEBUG - Final intent state: {state['intent']}")

    missing = [f for f in intent_fields if state["intent"].get(f) is None]
    print(f"DEBUG - Missing intent fields: {missing}")
    if missing:
        # Check if user is asking for properties, images, or address (using AI classification)
        classification = state.get("classification")
        print(f"DEBUG - Intent update classification check: '{classification}'")
        
        # Check if this is a request that should skip preference questions
        skip_preferences = (
            classification == "search_listings" or
            classification.startswith("request_images") or
            classification == "request_address" or
            classification == "request_details"
        )
        
        if skip_preferences:
            # User is asking for properties, images, or address - don't interrupt with preference questions
            print(f"DEBUG - User asking for {classification}, skipping preference questions")
            
            # Special case: if user is asking for properties but has no location, ask for location first
            if classification == "search_listings" and not state["intent"].get("location_preference"):
                print("DEBUG - User asking for properties but no location, asking for location")
                state["response"] = "I'd be happy to help you find properties! Which area or city are you interested in?"
                return state
            
            return {
                "response": None,
                "classification": classification
            }
        
        # If user is not actively searching for properties, don't ask preference questions
        # Let conversation flow naturally - only ask preferences when they're actually looking
        print(f"DEBUG - User not actively searching for properties (classification: {classification}), letting conversation flow naturally")
        return {
            "response": None,
            "classification": classification
        }

    print("DEBUG - Intent update returning None")
    return {
        "response": None,
        "classification": state.get("classification")
    }

async def classify_user_intent(state):
    """
    Classify the user's message to determine if they want to search for properties.
    """
    print("DEBUG - Starting classify_user_intent")
    user_message = state["user_message"]
    
    # Get available properties for context
    props = state.get("properties", [])
    prop_titles = [p.get("title", "").lower() for p in props]
    
    # Get recent conversation context
    session_messages = state.get("session_messages", [])
    recent_context = ""
    if session_messages:
        recent_messages = session_messages[-5:]  # Last 5 messages
        recent_context = "\n".join([f"{msg.get('role')}: {msg.get('content', '')}" for msg in recent_messages])
    
    prompt = f"""
Classify the user's message into exactly one of:
- search_listings  (user explicitly wants to see property listings)
- request_images   (user wants to see images/photos/pictures of a listing)
- request_address  (user wants the address/location of a listing)
- request_details  (user wants specific property info like price, features, floorplan, video, size, etc.)
- other             (anything else, including greetings, goodbyes, general conversation, end of conversation, casual location mentions)

IMPORTANT: Only classify as "search_listings" if the user is EXPLICITLY asking to see properties, listings, or available options.

Available properties for context: {prop_titles}

Recent conversation context:
{recent_context}

Property identifier extraction rules:
1. If user says "option X" or "option X images" → request_images:option X
2. If user mentions a property type (warehouse, office, space) and it matches a property title → request_images:PROPERTY_TYPE
3. If user says "this property", "that property", "the property" → request_images:this
4. If user says "show me images" without specific reference → request_images
5. If user mentions specific location/area that matches a property → request_images:LOCATION
6. If user asks for images in context where only one property is being discussed → request_images:this
7. If user uses pronouns like "it", "this", "that" when asking for images → request_images:this

Context analysis:
- If the user has been discussing a specific property and now asks for images, classify as request_images:this
- If there's only one property available and user asks for images, classify as request_images:this
- If user asks for images after a property search, assume they want images of the most relevant property

End of conversation indicators (classify as "other"):
- Goodbyes: "bye", "goodbye", "see you", "take care"
- Completion: "I'm all good", "that's all", "no thanks", "I'm done"
- General conversation: greetings, casual chat, non-property related topics
- Simple location mentions without asking for properties: "I'm in Johannesburg", "I live in Cape Town"

Examples of search_listings (user explicitly asking for properties):
- "What properties do you have in JHB?"
- "Show me available listings"
- "Do you have any properties for sale?"
- "Any properties available?"
- "Show me properties"
- "What do you have available?"
- "I'm looking for properties in Cape Town"
- "I want to rent a warehouse in Johannesburg"
- "Looking to buy office space in Sandton"
- "Do you have any rentals available?"
- "What's for sale in Cape Town?"

Examples of other (NOT property searches):
- "Hi" → other
- "Hello" → other
- "I'm in Johannesburg" → other
- "Johannesburg" → other
- "Cape Town" → other
- "How are you?" → other
- "I'm all good" → other
- "Goodbye" → other
- "Thanks for your help" → other
- "That's all I need" → other

Examples of request_details:
- "How much is this warehouse?" → request_details
- "What is the price?" → request_details  
- "What are the features?" → request_details
- "How big is it?" → request_details

Examples of request_images:
- "Show me images" → request_images
- "Show me images of option 1" → request_images:option 1
- "I want to see the warehouse images" → request_images:warehouse
- "Show me pictures of this property" → request_images:this
- "Can I see images of it?" → request_images:this
- "Show me the images" → request_images:this

Return only the tag (and identifier if applicable).
Message: {user_message}
"""
    resp = await llm.ainvoke([{"role":"user","content":prompt}])
    classification = resp.content.strip()
    state["classification"] = classification
    print(f"DEBUG - Classification result: '{classification}' for message: '{user_message}'")
    return {"classification": classification, "response": None}

async def extract_and_search_properties(state):
    """
    Search for properties based on user intent and store results in state.
    """
    print("DEBUG - Starting extract_and_search_properties")
    # Only search when the LLM tagged this as a listings request
    classification = state.get("classification")
    print(f"DEBUG - Property search classification check: '{classification}'")
    
    # Check if classification matches our search categories
    is_search_request = (
        classification == "search_listings" or
        classification.startswith("request_images") or
        classification == "request_address" or
        classification == "request_details"
    )
    
    if not is_search_request:
        print(f"DEBUG - Skipping property search, classification is '{classification}'")
        return {"response": None}

    intent = state.get("intent", {})
    user_message = state.get("user_message", "").lower()
    print(f"DEBUG - extract_and_search_properties intent: {intent}")
    print(f"DEBUG - User message: {user_message}")
    
    # Check if we have the minimum required field for property search
    location = intent.get("location_preference")
    
    if not location:
        # Missing location, but user is asking for properties - ask for location
        print("DEBUG - No location found, asking for location")
        state["response"] = "I'd be happy to help you find properties! Which area or city are you interested in?"
        return state
    
    # Prepare filters for property search
    filters = {"location_preference": location}
    
    # Add budget if set
    if intent.get("budget") is not None:
        filters["budget"] = intent["budget"]
    
    # Add size if set
    if intent.get("size_preference_sqm") is not None:
        filters["size_preference_sqm"] = intent["size_preference_sqm"]
    
    # Add transaction type if set (rent vs buy)
    if intent.get("transaction_type") is not None:
        filters["transaction_type"] = intent["transaction_type"]
    
    # Search for properties with flexible ranges
    print(f"DEBUG - Searching with filters: {filters}")
    properties = await search_properties(filters)
    state["properties"] = properties
    print(f"DEBUG - Found {len(properties)} properties with flexible ranges")
    
    if properties:
        print("DEBUG - Properties found, returning properties to continue to chat")
        return {
            "properties": properties,
            "classification": state.get("classification")
        }

    # No properties found with any filters
    state["response"] = (
        f"I don't have any listings right now in {location}. "
        "I'll notify you as soon as something becomes available. "
        "Feel free to reach out any time!"
    )
    await end_session(state["session_id"])
    return state

async def detect_end_chat(state):
    """
    Detect if the user wants to end the chat session using AI.
    """
    user_message = state["user_message"]
    session_id = state["session_id"]
    
    # Use AI to detect if user wants to end the conversation
    prompt = f"""
Analyze this user message and determine if they want to end the conversation or are saying goodbye.

User message: "{user_message}"

Consider these scenarios as "end conversation":
- User is satisfied and ending the conversation (e.g., "I'm all good", "that's all", "no thanks")
- User is saying goodbye (e.g., "bye", "goodbye", "see you", "take care")
- User is declining further assistance (e.g., "not interested", "maybe later", "I'll think about it")
- User is indicating they're done (e.g., "I'm done", "that's it", "nothing else")
- User is politely ending the conversation (e.g., "thanks for your help", "appreciate it")

Consider these scenarios as "continue conversation":
- User is asking questions about properties
- User is providing feedback but wants to continue
- User is making small talk but not ending
- User is asking for more information

Respond with ONLY:
- "end" if the user wants to end the conversation
- "continue" if the user wants to continue the conversation

Response:"""
    
    try:
        response = await llm.ainvoke([{"role": "user", "content": prompt}])
        result = response.content.strip().lower()
        
        print(f"DEBUG - End chat AI detection: '{result}' for message: '{user_message}'")
        
        if result == "end":
            await end_session(session_id)
            return {"response": "Thanks for chatting! I've ended this session. Goodbye!"}
        
        return {"response": None}
        
    except Exception as e:
        print(f"DEBUG - Error in end chat detection: {e}")
        # Fallback to continue conversation if AI fails
        return {"response": None}

# --- Build LangGraph ---
graph = StateGraph(ChatState)
graph.add_node("persona_update", RunnableLambda(extract_and_update_persona))
graph.add_node("exit_check_early", RunnableLambda(detect_end_chat))
graph.add_node("classify_intent", RunnableLambda(classify_user_intent))
graph.add_node("intent_update", RunnableLambda(extract_and_update_intent))
graph.add_node("property_search", RunnableLambda(extract_and_search_properties))
graph.add_node("exit_check", RunnableLambda(detect_end_chat))
graph.add_node("chat", RunnableLambda(chat_with_session_memory))
graph.set_entry_point("persona_update")

# Add conditional edges - if a node returns a response, go to END
def should_continue(state):
    """Check if we should continue to the next node or end"""
    has_response = state.get("response") is not None
    print(f"DEBUG - should_continue: response={state.get('response')}, has_response={has_response}, should_continue={not has_response}")
    return not has_response

graph.add_edge("persona_update", "exit_check_early")
graph.add_conditional_edges("exit_check_early", should_continue, {
    True: "classify_intent",
    False: END
})
graph.add_edge("classify_intent", "intent_update")
graph.add_conditional_edges("intent_update", should_continue, {
    True: "property_search",
    False: END
})
graph.add_conditional_edges("property_search", should_continue, {
    True: "exit_check",
    False: END
})
graph.add_conditional_edges("exit_check", should_continue, {
    True: "chat",
    False: END
})
graph.add_edge("chat", END)
chat_graph = graph.compile()

async def process_message(user_message: str, user_info: dict = None, session_id: str = None, wa_id: str = None, wamid: str = None, persona: dict = None, intent: dict = None, properties: list = None):
    """Process a message through the AI chat system with session memory"""
    if user_info is None:
        user_info = {}
    
    # Get session messages for context
    session_messages = []
    if session_id:
        session_messages = await get_session_messages(session_id, limit=10)
    
    # Process with AI
    result = await chat_graph.ainvoke({
        "user_message": user_message,
        "user_info": user_info,
        "session_id": session_id,
        "wa_id": wa_id,
        "wamid": wamid,
        "session_messages": session_messages,
        "persona": persona or {},
        "intent": intent or {},
        "properties": properties or []
    })
    
    # Check if this is an image request and handle it
    classification = result.get("classification", "")
    if classification.startswith("request_images"):
        print("DEBUG - Processing image request in process_message")
        # Create state for image request handling
        image_state = {
            "user_message": user_message,
            "properties": result.get("properties", []),
            "classification": classification,
            "session_messages": session_messages
        }
        
        # Handle image request
        image_messages = await handle_image_request(image_state)
        
        if image_messages:
            # Save the first text message to database
            if session_id and wa_id and wamid:
                await save_message(session_id, wa_id, wamid, "user", user_message)
                # Save the first text response
                if isinstance(image_messages[0], str):
                    await save_message(session_id, wa_id, f"{wamid}_ai", "assistant", image_messages[0])
            
            return {
                "response": image_messages[0] if isinstance(image_messages[0], str) else "Here are the images you requested!",
                "properties": result.get("properties", []),
                "classification": classification,
                "image_messages": image_messages  # Additional field for image handling
            }
    
    # Save messages to database
    if session_id and wa_id and wamid:
        await save_message(session_id, wa_id, wamid, "user", user_message)
        await save_message(session_id, wa_id, f"{wamid}_ai", "assistant", result["response"])
    
    return {
        "response": result["response"],
        "properties": result.get("properties", []),
        "classification": result.get("classification", "")
    }

async def handle_image_request(state):
    """
    Handle requests for property images and return image messages to send.
    """
    user_message = state["user_message"].lower()
    props = state.get("properties", [])
    classification = state.get("classification", "")
    
    print(f"DEBUG - handle_image_request: classification='{classification}', props count={len(props)}")
    
    # Check if this is an image request
    if not classification.startswith("request_images") or not props:
        print(f"DEBUG - Image request check failed: classification starts with request_images? {classification.startswith('request_images')}, has props? {len(props) > 0}")
        return None
    
    # Extract property identifier from classification if present
    property_identifier = None
    if ":" in classification:
        property_identifier = classification.split(":", 1)[1].lower()
    
    print(f"DEBUG - Property identifier: '{property_identifier}'")
    print(f"DEBUG - Available properties: {[p.get('title') for p in props]}")
    
    # Smart property selection based on AI classification
    selected_property = None
    
    if property_identifier:
        # Method 1: Handle option numbers
        if "option" in property_identifier:
            # Extract numeric option
            import re
            numbers = re.findall(r'\d+', property_identifier)
            if numbers:
                option_num = int(numbers[0])
                if 1 <= option_num <= len(props):
                    selected_property = props[option_num - 1]
        
        # Method 2: Handle text-based identifiers
        if not selected_property:
            best_match_score = 0
            for prop in props:
                title = prop.get("title", "").lower()
                location = prop.get("location", "").lower()
                city = prop.get("city", "").lower()
                
                # Check if identifier matches property keywords
                score = 0
                identifier_words = property_identifier.split()
                
                for word in identifier_words:
                    if word in title:
                        score += 3  # Title matches are most important
                    if word in location:
                        score += 2
                    if word in city:
                        score += 1
                    # Check for property type keywords
                    if word in ["office", "warehouse", "space"] and word in title:
                        score += 2
                
                if score > best_match_score:
                    best_match_score = score
                    selected_property = prop
    
    # Look for specific property selections in conversation
    session_messages = state.get("session_messages", [])
    recent_messages = session_messages[-20:]  # Look at more messages

    # Look for patterns like "option 3", "the warehouse", "this property"
    selected_property_context = None

    for msg in recent_messages:
        if msg.get("role") == "user":
            content = msg.get("content", "").lower()
            # Look for option selections
            if "option" in content:
                import re
                option_match = re.search(r'option\s+(\d+)', content)
                if option_match:
                    option_num = int(option_match.group(1))
                    if 1 <= option_num <= len(props):
                        selected_property_context = props[option_num - 1]
                        print(f"DEBUG - Found user selected option {option_num}: {selected_property_context.get('title')}")
                        break
            
            # Look for specific property type mentions that user explicitly asked about
            # Only if they specifically mentioned the property type in their image request
            if "warehouse" in content and "warehouse" in user_message:
                for i, prop in enumerate(props):
                    if "warehouse" in prop.get("title", "").lower():
                        selected_property_context = prop
                        print(f"DEBUG - Found user specifically asking for warehouse images: {prop.get('title')}")
                        break
            elif "office" in content and "office" in user_message:
                for i, prop in enumerate(props):
                    if "office" in prop.get("title", "").lower():
                        selected_property_context = prop
                        print(f"DEBUG - Found user specifically asking for office images: {prop.get('title')}")
                        break

    # Use the property the user specifically selected/discussed
    if selected_property_context:
        selected_property = selected_property_context
    
    # Enhanced context analysis for "this", "that", pronouns, etc.
    if not selected_property:
        # Check if user is using pronouns or demonstratives
        user_message_lower = state["user_message"].lower()
        is_pronoun_reference = any(word in user_message_lower for word in ["this", "that", "it", "the property", "the listing"])
        
        if is_pronoun_reference:
            print("DEBUG - User using pronoun/demonstrative reference, analyzing recent context")
            
            # Look for the most recently discussed property in the conversation
            recent_messages = session_messages[-10:]  # Last 10 messages
            property_mentions = {}
            
            for msg in recent_messages:
                content = msg.get("content", "").lower()
                for i, prop in enumerate(props):
                    title = prop.get("title", "").lower()
                    location = prop.get("location", "").lower()
                    city = prop.get("city", "").lower()
                    
                    # Check for property mentions with higher weight for recent messages
                    score = 0
                    title_words = title.split()
                    
                    # Check if property title words are mentioned
                    for word in title_words[:3]:  # First 3 words of title
                        if word in content:
                            score += 2
                    
                    # Check if location is mentioned
                    if location in content:
                        score += 1
                    
                    # Check if city is mentioned
                    if city in content:
                        score += 1
                    
                    # Give higher weight to more recent messages
                    message_index = recent_messages.index(msg)
                    recency_weight = 10 - message_index  # More recent = higher weight
                    score *= recency_weight
                    
                    if score > 0:
                        property_mentions[i] = property_mentions.get(i, 0) + score
                        print(f"DEBUG - Found mention of property {i} (score {score}): {title}")
            
            # Use most mentioned property from recent conversation
            if property_mentions:
                most_mentioned = max(property_mentions.items(), key=lambda x: x[1])
                selected_property = props[most_mentioned[0]]
                print(f"DEBUG - Selected property from pronoun context: {selected_property.get('title')}")
    
    # If still no property selected and user didn't specify, ask for clarification
    if not selected_property:
        print("DEBUG - No clear property context found, asking user to specify")
        prop_options = []
        for i, prop in enumerate(props[:3], 1):  # Show first 3 options
            prop_options.append(f"Option {i}: {prop.get('title')}")
        
        options_text = "\n".join(prop_options)
        return [f"I have multiple properties available. Which one would you like to see images of?\n\n{options_text}\n\nPlease let me know which option you'd like images for."]
    
    # Final fallback: use first property if only one or no context found
    if not selected_property:
        selected_property = props[0]
    
    print(f"DEBUG - Selected property: '{selected_property.get('title')}'")
    
    # Get images from selected property
    images = selected_property.get("images", [])
    
    print(f"DEBUG - Images found: {len(images) if images else 0}")
    print(f"DEBUG - Image URLs: {images}")
    
    if not images:
        property_title = selected_property.get("title", "this listing")
        return [f"Sorry, I don't have any images available for {property_title}."]
    
    # Prepare image messages
    image_messages = []
    property_title = selected_property.get("title", "This property")
    
    # Add a text message first to introduce the images
    image_messages.append(f"Here are the images for {property_title}:")
    
    for i, image_url in enumerate(images[:5]):  # Limit to 5 images
        caption = f"{property_title} - Image {i+1}" if i > 0 else f"{property_title}"
        image_messages.append({
            "type": "image",
            "url": image_url,
            "caption": caption
        })
    
    return image_messages