ai_chat.py

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

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", "")
    

    
    # 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."
    )
    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", {})
    must_have_list = intent_data.get('must_have', []) or []
    system_message += (
        f" They're looking for a property in {intent_data.get('location_preference','[any area]')}, "
        f"with a budget up to {intent_data.get('budget','[any amount]')} per month, "
        f"around {intent_data.get('size_preference_sqm','[size]')} sqm, "
        f"and must-haves: {', '.join(must_have_list) if must_have_list else '[none]'}. "
    )
    
    # 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 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"
            if p.get("features"):
                system_message += f"  Features: {', '.join(p.get('features', [])[:5])}\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 available_extras:
                system_message += f"  Available on request: {', '.join(available_extras)}\n"
    
    system_message += "\n\nIMPORTANT: You may only reference listings passed in state['properties']. When users ask for images, photos, or pictures, let them know that images are available and will be sent separately. When users ask for the address, location, or where a property is located, provide the Google Maps address from the property data. The addresses are Google Maps compatible for navigation. If information is not available, respond with 'For this listing, I don't have [specific detail] available right now'."
    
    # Build messages array with history
    messages = [{"role": "system", "content": system_message}]
    
    # Add conversation history (last 10 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
        
        # 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", {})

    # 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.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

    # e. If any field still unset, ask a gentle follow-up
    missing = [f for f in persona_fields if state["persona"].get(f) is None]
    if missing:
        state["response"] = f"Hi there! What is your {missing[0]} preference?"
        print(f"DEBUG - Persona update returning response: {state['response']}")
        return state

    # f. All persona fields present—proceed to chat
    print("DEBUG - Persona update returning None")
    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"]
    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', [])}

    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. 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.
    4. 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
            }
        
        # User is setting preferences, ask for missing fields
        questions = {
            "location_preference": "Hi there! Which area or suburb are you interested in?",
            "budget": "Hi there! What is your monthly budget?",
            "size_preference_sqm": "Hi there! How many square metres do you need?",
            "must_have": "Hi there! What features are must-haves for you?"
        }
        state["response"] = questions.get(missing[0], f"Hi there! Could you tell me your {missing[0]}?")
        print(f"DEBUG - Intent update returning response: {state['response']}")
        return state

    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"]
    prompt = f"""
Classify the user's message into exactly one of:
- search_listings  (user 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)

If the user is asking for images and mentions a specific property (like "option 1", "the office", "warehouse", etc.), 
extract the property identifier and return: request_images:IDENTIFIER

Examples:
- "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
- "Show me images" → request_images
- "What do you have in JHB?" → search_listings
- "Do you have any properties for sale?" → search_listings
- "Any properties available?" → search_listings
- "Show me properties" → search_listings

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"]
    
    # 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.
    """
    user_message = state["user_message"].lower()
    session_id = state["session_id"]
    
    if any(phrase in user_message for phrase in ["thank you", "thanks", "bye", "goodbye", "end chat"]):
        await end_session(session_id)
        return {"response": "Thanks for chatting! I've ended this session. Goodbye!"}
    
    return {"response": None}

# --- Build LangGraph ---
graph = StateGraph(ChatState)
graph.add_node("persona_update", RunnableLambda(extract_and_update_persona))
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", "classify_intent")
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 []
    })
    
    # 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 property type mentions that user specifically asked about
            for i, prop in enumerate(props):
                title_words = prop.get("title", "").lower().split()
                for word in ["warehouse", "office", "space", "unit"]:
                    if word in content and word in title_words:
                        selected_property_context = prop
                        print(f"DEBUG - Found user interest in {word}: {prop.get('title')}")
                        break

    # Use the property the user specifically selected/discussed
    if selected_property_context:
        selected_property = selected_property_context
    
    # Fallback: Use conversation context to find which property user was discussing
    if not selected_property and len(props) > 1:
        # Check conversation history for property context
        session_messages = state.get("session_messages", [])
        print(f"DEBUG - Checking {len(session_messages)} session messages for property context")
        
        # Look for property mentions in recent conversation
        recent_messages = session_messages[-10:]  # Last 10 messages
        property_mentions = {}
        
        for msg in recent_messages:
            if msg.get("role") == "assistant":
                content = msg.get("content", "").lower()
                for i, prop in enumerate(props):
                    title = prop.get("title", "").lower()
                    location = prop.get("location", "").lower()
                    # Check if this property was mentioned in AI response
                    title_words = title.split()
                    if len(title_words) >= 2:  # Use first 2 words for matching
                        key_phrase = " ".join(title_words[:2])
                        if key_phrase in content or location in content:
                            property_mentions[i] = property_mentions.get(i, 0) + 1
                            print(f"DEBUG - Found mention of property {i}: {title}")
        
        # Use most mentioned property from 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 conversation context: {selected_property.get('title')}")
        else:
            # Multiple properties available - ask 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