Update app.py

app.py

@@ -25,10 +25,9 @@ from sqlalchemy import Column, Integer, String, DateTime, Text, Float
 # --- Environment Variables and API Keys ---
 SPOONACULAR_API_KEY = os.getenv("SPOONACULAR_API_KEY", "default_fallback_value")
 PAYSTACK_SECRET_KEY = os.getenv("PAYSTACK_SECRET_KEY", "default_fallback_value")
-DATABASE_URL = os.getenv("DATABASE_URL", "default_fallback_value")  # Example using SQLite
+DATABASE_URL = os.getenv("DATABASE_URL", "default_fallback_value")  # Example using SQLite
 NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY", "default_fallback_value")
 openai.api_key = os.getenv("OPENAI_API_KEY", "default_fallback_value")
-GOOGLE_MAPS_API_KEY = os.getenv("GOOGLE_MAPS_API_KEY", "default_fallback_value")  # Add Google Maps API key
 
 # WhatsApp Business API credentials (Cloud API)
 WHATSAPP_PHONE_NUMBER_ID = os.getenv("WHATSAPP_PHONE_NUMBER_ID", "default_value")
@@ -37,935 +36,942 @@ MANAGEMENT_WHATSAPP_NUMBER = os.getenv("MANAGEMENT_WHATSAPP_NUMBER", "default_va
 
 # Synthetic town prices for shipping costs
 TOWN_SHIPPING_COSTS = {
-    "lasu gate": 1000,  # N1,000 for LASU Gate
-    "ojo": 800,         # N800 for Ojo
-    "ajangbadi": 1200,  # N1,200 for Ajangbadi
-    "iba": 900,         # N900 for Iba
-    "okokomaiko": 1500, # N1,500 for Okokomaiko
-    "default": 1000     # Default shipping cost for unknown areas
+    "lasu gate": 1000,  # N1,000 for LASU Gate
+    "ojo": 800,         # N800 for Ojo
+    "ajangbadi": 1200,  # N1,200 for Ajangbadi
+    "iba": 900,         # N900 for Iba
+    "okokomaiko": 1500, # N1,500 for Okokomaiko
+    "default": 1000     # Default shipping cost for unknown areas
 }
 
 # --- Database Setup ---
 Base = declarative_base()
 
 class ChatHistory(Base):
-    __tablename__ = "chat_history"
-    id = Column(Integer, primary_key=True, index=True)
-    user_id = Column(String, index=True)
-    timestamp = Column(DateTime, default=datetime.utcnow)
-    direction = Column(String)  # 'inbound' or 'outbound'
-    message = Column(Text)
+    __tablename__ = "chat_history"
+    id = Column(Integer, primary_key=True, index=True)
+    user_id = Column(String, index=True)
+    timestamp = Column(DateTime, default=datetime.utcnow)
+    direction = Column(String)  # 'inbound' or 'outbound'
+    message = Column(Text)
 
 class Order(Base):
-    __tablename__ = "orders"
-    id = Column(Integer, primary_key=True, index=True)
-    order_id = Column(String, unique=True, index=True)
-    user_id = Column(String, index=True)
-    dish = Column(String)
-    quantity = Column(String)
-    price = Column(String, default="0")
-    status = Column(String, default="Pending Payment")
-    payment_reference = Column(String, nullable=True)
-    delivery_address = Column(String, default="")  # New field for address
-    timestamp = Column(DateTime, default=datetime.utcnow)
+    __tablename__ = "orders"
+    id = Column(Integer, primary_key=True, index=True)
+    order_id = Column(String, unique=True, index=True)
+    user_id = Column(String, index=True)
+    dish = Column(String)
+    quantity = Column(String)
+    price = Column(String, default="0")
+    status = Column(String, default="Pending Payment")
+    payment_reference = Column(String, nullable=True)
+    delivery_address = Column(String, default="")  # New field for address
+    timestamp = Column(DateTime, default=datetime.utcnow)
 
 class UserProfile(Base):
-    __tablename__ = "user_profiles"
-    id = Column(Integer, primary_key=True, index=True)
-    user_id = Column(String, unique=True, index=True)
-    phone_number = Column(String, unique=True, index=True, nullable=True)
-    name = Column(String, default="Valued Customer")
-    email = Column(String, default="unknown@example.com")
-    preferences = Column(Text, default="")
-    last_interaction = Column(DateTime, default=datetime.utcnow)
-    order_ids = Column(Text, default="")  # Store order IDs as a comma-separated string
+    __tablename__ = "user_profiles"
+    id = Column(Integer, primary_key=True, index=True)
+    user_id = Column(String, unique=True, index=True)
+    phone_number = Column(String, unique=True, index=True, nullable=True)
+    name = Column(String, default="Valued Customer")
+    email = Column(String, default="unknown@example.com")
+    preferences = Column(Text, default="")
+    last_interaction = Column(DateTime, default=datetime.utcnow)
 
 class SentimentLog(Base):
-    __tablename__ = "sentiment_logs"
-    id = Column(Integer, primary_key=True, index=True)
-    user_id = Column(String, index=True)
-    timestamp = Column(DateTime, default=datetime.utcnow)
-    sentiment_score = Column(Float)
-    message = Column(Text)
-
+    __tablename__ = "sentiment_logs"
+    id = Column(Integer, primary_key=True, index=True)
+    user_id = Column(String, index=True)
+    timestamp = Column(DateTime, default=datetime.utcnow)
+    sentiment_score = Column(Float)
+    message = Column(Text)
+
+# --- New Model: Order Tracking ---
 class OrderTracking(Base):
-    __tablename__ = "order_tracking"
-    id = Column(Integer, primary_key=True, index=True)
-    order_id = Column(String, index=True)
-    status = Column(String)  # e.g., "Order Placed", "Payment Confirmed", etc.
-    message = Column(Text, nullable=True)  # Optional additional details
-    timestamp = Column(DateTime, default=datetime.utcnow)
+    __tablename__ = "order_tracking"
+    id = Column(Integer, primary_key=True, index=True)
+    order_id = Column(String, index=True)
+    status = Column(String)  # e.g., "Order Placed", "Payment Confirmed", etc.
+    message = Column(Text, nullable=True)  # Optional additional details
+    timestamp = Column(DateTime, default=datetime.utcnow)
 
 # --- Create Engine and Session ---
 engine = create_async_engine(DATABASE_URL, echo=True)
 async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
 
 async def init_db():
-    async with engine.begin() as conn:
-        await conn.run_sync(Base.metadata.create_all)
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
 
 # --- Global In-Memory Stores ---
-user_state = {}       # e.g., { user_id: ConversationState }
-conversation_context = {}  # { user_id: [ { "timestamp": ..., "role": "user"/"bot", "message": ... }, ... ] }
+user_state = {}       # e.g., { user_id: ConversationState }
+conversation_context = {}  # { user_id: [ { "timestamp": ..., "role": "user"/"bot", "message": ... }, ... ] }
 proactive_timer = {}
 
 menu_items = [
-    {"name": "Jollof Rice", "description": "A spicy and flavorful rice dish", "price": 1500, "nutrition": "Calories: 300 kcal, Carbs: 50g, Protein: 10g, Fat: 5g"},
-    {"name": "Fried Rice", "description": "A savory rice dish with vegetables and meat", "price": 1200, "nutrition": "Calories: 350 kcal, Carbs: 55g, Protein: 12g, Fat: 8g"},
-    {"name": "Chicken Wings", "description": "Crispy fried chicken wings", "price": 2000, "nutrition": "Calories: 400 kcal, Carbs: 20g, Protein: 25g, Fat: 15g"},
-    {"name": "Egusi Soup", "description": "A rich and hearty soup made with melon seeds", "price": 1000, "nutrition": "Calories: 250 kcal, Carbs: 15g, Protein: 8g, Fat: 10g"}
+    {"name": "Jollof Rice", "description": "A spicy and flavorful rice dish", "price": 1500, "nutrition": "Calories: 300 kcal, Carbs: 50g, Protein: 10g, Fat: 5g"},
+    {"name": "Fried Rice", "description": "A savory rice dish with vegetables and meat", "price": 1200, "nutrition": "Calories: 350 kcal, Carbs: 55g, Protein: 12g, Fat: 8g"},
+    {"name": "Chicken Wings", "description": "Crispy fried chicken wings", "price": 2000, "nutrition": "Calories: 400 kcal, Carbs: 20g, Protein: 25g, Fat: 15g"},
+    {"name": "Egusi Soup", "description": "A rich and hearty soup made with melon seeds", "price": 1000, "nutrition": "Calories: 250 kcal, Carbs: 15g, Protein: 8g, Fat: 10g"}
 ]
 
 # --- Conversation State Management ---
 SESSION_TIMEOUT = timedelta(minutes=5)
 
 class ConversationState:
-    def __init__(self):
-        self.flow = None  # e.g., "order"
-        self.step = 0
-        self.data = {}
-        self.last_active = datetime.utcnow()
+    def __init__(self):
+        self.flow = None  # e.g., "order"
+        self.step = 0
+        self.data = {}
+        self.last_active = datetime.utcnow()
 
-    def update_last_active(self):
-        self.last_active = datetime.utcnow()
+    def update_last_active(self):
+        self.last_active = datetime.utcnow()
 
-    def is_expired(self):
-        return datetime.utcnow() - self.last_active > SESSION_TIMEOUT
+    def is_expired(self):
+        return datetime.utcnow() - self.last_active > SESSION_TIMEOUT
 
-    def reset(self):
-        self.flow = None
-        self.step = 0
-        self.data = {}
-        self.last_active = datetime.utcnow()
+    def reset(self):
+        self.flow = None
+        self.step = 0
+        self.data = {}
+        self.last_active = datetime.utcnow()
 
 # --- Utility Functions ---
 async def log_chat_to_db(user_id: str, direction: str, message: str):
-    async with async_session() as session:
-        entry = ChatHistory(user_id=user_id, direction=direction, message=message)
-        session.add(entry)
-        await session.commit()
+    async with async_session() as session:
+        entry = ChatHistory(user_id=user_id, direction=direction, message=message)
+        session.add(entry)
+        await session.commit()
 
 async def log_sentiment(user_id: str, message: str, score: float):
-    async with async_session() as session:
-        entry = SentimentLog(user_id=user_id, sentiment_score=score, message=message)
-        session.add(entry)
-        await session.commit()
+    async with async_session() as session:
+        entry = SentimentLog(user_id=user_id, sentiment_score=score, message=message)
+        session.add(entry)
+        await session.commit()
 
 def analyze_sentiment(text: str) -> float:
-    blob = TextBlob(text)
-    return blob.sentiment.polarity
+    blob = TextBlob(text)
+    return blob.sentiment.polarity
 
 def google_image_scrape(query: str) -> str:
-    headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"}
-    search_url = f"https://www.google.com/search?tbm=isch&q={query}"
-    try:
-        response = requests.get(search_url, headers=headers, timeout=5)
-    except Exception:
-        return ""
-    if response.status_code == 200:
-        soup = BeautifulSoup(response.text, "html.parser")
-        img_tags = soup.find_all("img")
-        for img in img_tags:
-            src = img.get("src")
-            if src and src.startswith("http"):
-                return src
-    return ""
+    headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"}
+    search_url = f"https://www.google.com/search?tbm=isch&q={query}"
+    try:
+        response = requests.get(search_url, headers=headers, timeout=5)
+    except Exception:
+        return ""
+    if response.status_code == 200:
+        soup = BeautifulSoup(response.text, "html.parser")
+        img_tags = soup.find_all("img")
+        for img in img_tags:
+            src = img.get("src")
+            if src and src.startswith("http"):
+                return src
+    return ""
 
 def create_paystack_payment_link(email: str, amount: int, reference: str) -> dict:
-    url = "https://api.paystack.co/transaction/initialize"
-    headers = {
-        "Authorization": f"Bearer {PAYSTACK_SECRET_KEY}",
-        "Content-Type": "application/json",
-    }
-    data = {
-        "email": email,
-        "amount": amount,
-        "reference": reference,
-        "callback_url": "https://custy-bot.vercel.app/payment_callback"
-    }
-    try:
-        response = requests.post(url, json=data, headers=headers, timeout=10)
-        if response.status_code == 200:
-            return response.json()
-        else:
-            return {"status": False, "message": "Failed to initialize payment."}
-    except Exception as e:
-        return {"status": False, "message": str(e)}
+    url = "https://api.paystack.co/transaction/initialize"
+    headers = {
+        "Authorization": f"Bearer {PAYSTACK_SECRET_KEY}",
+        "Content-Type": "application/json",
+    }
+    data = {
+        "email": email,
+        "amount": amount,
+        "reference": reference,
+        "callback_url": "https://custy-bot.vercel.app/payment_callback"
+    }
+    try:
+        response = requests.post(url, json=data, headers=headers, timeout=10)
+        if response.status_code == 200:
+            return response.json()
+        else:
+            return {"status": False, "message": "Failed to initialize payment."}
+    except Exception as e:
+        return {"status": False, "message": str(e)}
 
 # --- WhatsApp Business API Helper ---
 def send_whatsapp_message(recipient: str, message_body: str) -> dict:
-    """
-    Sends a WhatsApp text message using the WhatsApp Cloud API.
-    `recipient` should be in international format, e.g., "15551234567".
-    """
-    url = f"https://graph.facebook.com/v15.0/{WHATSAPP_PHONE_NUMBER_ID}/messages"
-    headers = {
-        "Authorization": f"Bearer {WHATSAPP_ACCESS_TOKEN}",
-        "Content-Type": "application/json"
-    }
-    payload = {
-        "messaging_product": "whatsapp",
-        "to": recipient,
-        "type": "text",
-        "text": {"body": message_body}
-    }
-    response = requests.post(url, headers=headers, json=payload)
-    return response.json()
-
-def stream_text_completion(prompt: str):
-    from openai import OpenAI
-    client = OpenAI(
-        base_url="https://integrate.api.nvidia.com/v1",
-        api_key=NVIDIA_API_KEY
-    )
-    print(f"Using NVIDIA API Key: {NVIDIA_API_KEY}")  # Debugging
-
-    try:
-        completion = client.chat.completions.create(
-            model="meta/llama-3.1-405b-instruct",
-            messages=[{"role": "user", "content": prompt}],
-            temperature=0.2,
-            top_p=0.7,
-            max_tokens=1024,
-            stream=True
-        )
-        for chunk in completion:
-            if chunk.choices[0].delta.content is not None:
-                yield chunk.choices[0].delta.content
-    except Exception as e:
-        yield f"Error: {str(e)}"  # Handle errors gracefully
-
+    """
+    Sends a WhatsApp text message using the WhatsApp Cloud API.
+    `recipient` should be in international format, e.g., "15551234567".
+    """
+    url = f"https://graph.facebook.com/v15.0/{WHATSAPP_PHONE_NUMBER_ID}/messages"
+    headers = {
+        "Authorization": f"Bearer {WHATSAPP_ACCESS_TOKEN}",
+        "Content-Type": "application/json"
+    }
+    payload = {
+        "messaging_product": "whatsapp",
+        "to": recipient,
+        "type": "text",
+        "text": {"body": message_body}
+    }
+    response = requests.post(url, headers=headers, json=payload)
+    return response.json()
+
+# --- NVIDIA LLM Streaming Functions ---
 def stream_text_completion(prompt: str):
-    from openai import OpenAI
-    client = OpenAI(
-        base_url="https://integrate.api.nvidia.com/v1",
-        api_key=NVIDIA_API_KEY
-    )
-    print(f"Using NVIDIA API Key: {NVIDIA_API_KEY}")  # Debugging
-
-    try:
-        completion = client.chat.completions.create(
-            model="meta/llama-3.1-405b-instruct",
-            messages=[{"role": "user", "content": prompt}],
-            temperature=0.2,
-            top_p=0.7,
-            max_tokens=1024,
-            stream=True
-        )
-        for chunk in completion:
-            if chunk.choices[0].delta.content is not None:
-                yield chunk.choices[0].delta.content
-    except Exception as e:
-        yield f"Error: {str(e)}"  # Handle errors gracefully
+    from openai import OpenAI
+    client = OpenAI(
+        base_url="https://integrate.api.nvidia.com/v1",
+        api_key=NVIDIA_API_KEY
+    )
+    completion = client.chat.completions.create(
+        model="meta/llama-3.1-405b-instruct",
+        messages=[{"role": "user", "content": prompt}],
+        temperature=0.2,
+        top_p=0.7,
+        max_tokens=1024,
+        stream=True
+    )
+    for chunk in completion:
+        if chunk.choices[0].delta.content is not None:
+            yield chunk.choices[0].delta.content
+
+def stream_image_completion(image_b64: str):
+    invoke_url = "https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-90b-vision-instruct/chat/completions"
+    headers = {
+        "Authorization": f"Bearer {NVIDIA_API_KEY}",
+        "Accept": "text/event-stream"
+    }
+    payload = {
+        "model": "meta/llama-3.2-90b-vision-instruct",
+        "messages": [
+            {
+                "role": "user",
+                "content": f'What is in this image? <img src="data:image/png;base64,{image_b64}" />'
+            }
+        ],
+        "max_tokens": 512,
+        "temperature": 1.00,
+        "top_p": 1.00,
+        "stream": True
+    }
+    response = requests.post(invoke_url, headers=headers, json=payload, stream=True)
+    for line in response.iter_lines():
+        if line:
+            yield line.decode("utf-8") + "\n"
 
 # --- Helper Function for Order Tracking ---
 async def log_order_tracking(order_id: str, status: str, message: str = None):
-    async with async_session() as session:
-        tracking_entry = OrderTracking(
-            order_id=order_id,
-            status=status,
-            message=message
-        )
-        session.add(tracking_entry)
-        await session.commit()
+    async with async_session() as session:
+        tracking_entry = OrderTracking(
+            order_id=order_id,
+            status=status,
+            message=message
+        )
+        session.add(tracking_entry)
+        await session.commit()
 
 # --- Advanced Internal Flow: Order Processing & Payment Integration ---
 def calculate_shipping_cost(address: str) -> int:
-    """
-    Calculate shipping cost based on the user's address.
-    """
-    address_lower = address.lower()
-    for area, cost in TOWN_SHIPPING_COSTS.items():
-        if area in address_lower:
-            return cost
-    return TOWN_SHIPPING_COSTS["default"]  # Default shipping cost for unknown areas
-
-# --- Google Maps ETA Calculation ---
-def calculate_eta(destination: str) -> str:
-    """
-    Calculate ETA from the restaurant address to the customer's delivery address using Google Maps API.
-    """
-    if not GOOGLE_MAPS_API_KEY:
-        return "ETA unavailable (Google Maps API key missing)."
-
-    origin = "Plot 13 Isashi Road, Iyana Isashi, Off Lagos - Badagry Expy, Lagos"
-    url = f"https://maps.googleapis.com/maps/api/directions/json?origin={origin}&destination={destination}&key={GOOGLE_MAPS_API_KEY}"
-    
-    try:
-        response = requests.get(url, timeout=10)
-        if response.status_code == 200:
-            data = response.json()
-            if data.get("routes"):
-                duration = data["routes"][0]["legs"][0]["duration"]["text"]
-                return f"Estimated delivery time: {duration}"
-            else:
-                return "ETA unavailable (no route found)."
-        else:
-            return "ETA unavailable (API error)."
-    except Exception as e:
-        return f"ETA unavailable (error: {str(e)})."
-
-# --- Contextual Order Intent Detection ---
-def is_order_intent(message: str) -> bool:
-    """
-    Check if the user's message indicates an intent to place an order.
-    """
-    order_keywords = ["order", "menu", "dish", "food", "deliver", "hungry"]
-    order_phrases = ["I want to order", "Can I order", "I'd like to order", "get food", "place an order"]
-
-    message_lower = message.lower()
-    for phrase in order_phrases:
-        if phrase in message_lower:
-            return True
-    for keyword in order_keywords:
-        if keyword in message_lower:
-            # Ensure the keyword is not part of another word (e.g., "border")
-            if re.search(rf"\b{keyword}\b", message_lower):
-                return True
-    return False
-
-    
-
-# ... (previous code remains unchanged)
-
-async def track_order(user_id: str, order_id: str) -> str:
-    """
-    Track an order and return its status and ETA.
-    """
-    async with async_session() as session:
-        # Fetch order details
-        order_result = await session.execute(
-            select(Order).where(Order.order_id == order_id)
-        )
-        order = order_result.scalars().first()
-        if not order:
-            return "Order not found. Please check your order ID."
-
-        # Fetch tracking updates
-        tracking_result = await session.execute(
-            select(OrderTracking)
-            .where(OrderTracking.order_id == order_id)
-            .order_by(OrderTracking.timestamp)
-        )
-        tracking_updates = tracking_result.scalars().all()
-
-        # Calculate ETA
-        eta = calculate_eta(order.delivery_address)
-
-        # Prepare response
-        response = f"Order ID: {order_id}\nStatus: {order.status}\n"
-        if tracking_updates:
-            response += "Tracking Updates:\n"
-            for update in tracking_updates:
-                response += f"- {update.status} ({update.timestamp}): {update.message or 'No details'}\n"
-        response += f"\n{eta}"
-        return response
-
-# ... (rest of the code remains unchanged)
-
-# --- Update User Profile with Order IDs ---
-async def update_user_profile_with_order(user_id: str, order_id: str):
-    async with async_session() as session:
-        result = await session.execute(
-            select(UserProfile).where(UserProfile.user_id == user_id)
-        )
-        profile = result.scalars().first()
-        if profile:
-            if profile.order_ids:
-                profile.order_ids += f",{order_id}"
-            else:
-                profile.order_ids = order_id
-            await session.commit()
-
-async def process_order_flow(user_id: str, message: str) -> str:
-    """
-    Implements an FSM-based order flow with shipping cost calculation.
-    """
-    # Retrieve or initialize conversation state
-    state = user_state.get(user_id)
-    if state and state.is_expired():
-        state.reset()
-        state.step = 1
-        state.update_last_active()
-        user_state[user_id] = state
-        if message.lower() == "order":
-            return "Sure! What dish would you like to order?"
-        return ""
-
-    # If no state exists but the message includes "order", start the order flow.
-    if not state and "order" in message.lower():
-        state = ConversationState()
-        state.flow = "order"
-        state.step = 1
-        state.update_last_active()
-        user_state[user_id] = state
-        return "Sure! What dish would you like to order?"
-
-    # --- New Logic: Parse Dish, Quantity, Phone Number, and Address in a Single Message ---
-    if not state or state.flow != "order":
-        # Check if the message contains a dish and quantity
-        dish_candidates = [item["name"] for item in menu_items]
-        found_dish = None
-        for dish in dish_candidates:
-            if dish.lower() in message.lower():
-                found_dish = dish
-                break
-        numbers = re.findall(r'\d+', message)
-        if found_dish and numbers:
-            quantity = int(numbers[0])
-            if quantity <= 0:
-                return "Please enter a valid quantity (e.g., 1, 2, 3)."
-            # Initialize state and skip to step 3 (phone number)
-            state = ConversationState()
-            state.flow = "order"
-            state.step = 3
-            state.data["dish"] = found_dish
-            state.data["quantity"] = quantity
-            state.update_last_active()
-            user_state[user_id] = state
-
-            # Extract phone number and address from the message
-            phone_pattern = r'(\+?\d{10,15})'
-            phone_match = re.search(phone_pattern, message)
-            address = None
-            if phone_match:
-                phone_number = phone_match.group(1)
-                # Assume the address starts after the phone number
-                address_start = phone_match.end()
-                address = message[address_start:].strip()
-                # Remove any leading/trailing commas or spaces
-                address = re.sub(r'^[,\s]+|[,\s]+$', '', address)
-            
-            if phone_match and address:
-                state.data["phone_number"] = phone_number
-                state.data["address"] = address
-                # Save phone number and address to the user's profile
-                await update_user_profile(user_id, phone_number, address)
-                # Calculate shipping cost based on the address
-                shipping_cost = calculate_shipping_cost(address)
-                state.data["shipping_cost"] = shipping_cost
-                state.step = 5
-                return (f"Thanks! Your phone number is recorded as: {phone_number}.\n"
-                        f"Your delivery address is: {address}.\n"
-                        f"Your delivery cost is N{shipping_cost}. Would you like to add any extras such as sides or drinks? (yes/no)")
-            elif phone_match:
-                state.data["phone_number"] = phone_match.group(1)
-                # Save phone number to the user's profile
-                await update_user_profile(user_id, phone_number)
-                return "Thank you. Please provide your delivery address."
-            else:
-                return "Please provide both your phone number and delivery address. For example: '09162409591, 1, Iyana Isashi, Isashi, Ojo, Lagos'."
-
-    if state and state.flow == "order":
-        state.update_last_active()
-        # --- Step 1: Expecting Dish Selection (and optionally quantity) ---
-        if state.step == 1:
-            dish_candidates = [item["name"] for item in menu_items]
-            found_dish = None
-            for dish in dish_candidates:
-                if dish.lower() in message.lower():
-                    found_dish = dish
-                    break
-            numbers = re.findall(r'\d+', message)
-            if found_dish:
-                state.data["dish"] = found_dish
-                if numbers:
-                    quantity = int(numbers[0])
-                    if quantity <= 0:
-                        return "Please enter a valid quantity (e.g., 1, 2, 3)."
-                    state.data["quantity"] = quantity
-                    state.step = 3  # Move to phone number step
-                    return f"You selected {found_dish} with {quantity} serving(s). Please provide your phone number and delivery address."
-                else:
-                    state.step = 2  # Ask for quantity
-                    return f"You selected {found_dish}. How many servings would you like?"
-            else:
-                return "I couldn't identify the dish. Please type the dish name from our menu."
-
-        # --- Step 2: Asking for Quantity ---
-        if state.step == 2:
-            numbers = re.findall(r'\d+', message)
-            if not numbers:
-                return "Please enter a valid number for the quantity (e.g., 1, 2, 3)."
-            quantity = int(numbers[0])
-            if quantity <= 0:
-                return "Please enter a valid quantity (e.g., 1, 2, 3)."
-            state.data["quantity"] = quantity
-            state.step = 3
-            return f"Got it. {quantity} serving(s) of {state.data.get('dish')}. Please provide your phone number and delivery address."
-
-        # --- Step 3: Parse Phone Number and Address from Message ---
-        if state.step == 3:
-            # Extract phone number using regex
-            phone_pattern = r'(\+?\d{10,15})'
-            phone_match = re.search(phone_pattern, message)
-            
-            # Extract address (assume it follows the phone number, separated by a comma)
-            address = None
-            if phone_match:
-                phone_number = phone_match.group(1)
-                # Assume the address starts after the phone number
-                address_start = phone_match.end()
-                address = message[address_start:].strip()
-                # Remove any leading/trailing commas or spaces
-                address = re.sub(r'^[,\s]+|[,\s]+$', '', address)
-            
-            if phone_match and address:
-                state.data["phone_number"] = phone_number
-                state.data["address"] = address
-                # Save phone number and address to the user's profile
-                await update_user_profile(user_id, phone_number, address)
-                # Calculate shipping cost based on the address
-                shipping_cost = calculate_shipping_cost(address)
-                state.data["shipping_cost"] = shipping_cost
-                state.step = 5
-                return (f"Thanks! Your phone number is recorded as: {phone_number}.\n"
-                        f"Your delivery address is: {address}.\n"
-                        f"Your delivery cost is N{shipping_cost}. Would you like to add any extras such as sides or drinks? (yes/no)")
-            elif phone_match:
-                state.data["phone_number"] = phone_match.group(1)
-                # Save phone number to the user's profile
-                await update_user_profile(user_id, phone_number)
-                return "Thank you. Please provide your delivery address."
-            else:
-                return "Please provide both your phone number and delivery address. For example: '09162409591, 1, Iyana Isashi, Isashi, Ojo, Lagos'."
-
-        # --- Step 4: Requesting Delivery Address ---
-        if state.step == 4:
-            state.data["address"] = message
-            # Save address to the user's profile
-            await update_user_profile(user_id, address=message)
-            # Calculate shipping cost based on the address
-            shipping_cost = calculate_shipping_cost(message)
-            state.data["shipping_cost"] = shipping_cost
-            state.step = 5
-            return (f"Thanks. Your delivery address is recorded as: {message}.\n"
-                    f"Your delivery cost is N{shipping_cost}. Would you like to add any extras such as sides or drinks? (yes/no)")
-
-        # --- Step 5: Asking if They Want Extras ---
-        if state.step == 5:
-            if message.lower() in ["yes", "y"]:
-                state.step = 6
-                return "Please list the extras you would like to add (e.g., drinks, sides, etc.)."
-            elif message.lower() in ["no", "n"]:
-                state.data["extras"] = ""
-                state.step = 7
-                dish = state.data.get("dish", "")
-                quantity = state.data.get("quantity", 1)
-                phone = state.data.get("phone_number", "")
-                address = state.data.get("address", "")
-                shipping_cost = state.data.get("shipping_cost", 0)
-                price_per_serving = 1500  # Fixed price per serving
-                total_price = (quantity * price_per_serving) + shipping_cost
-                summary = (f"Order Summary:\nDish: {dish}\nQuantity: {quantity}\n"
-                           f"Phone: {phone}\nAddress: {address}\n"
-                           f"Shipping Cost: N{shipping_cost}\n"
-                           f"Total Price: N{total_price}\n"
-                           f"Extras: None\nConfirm order? (yes/no)")
-                return summary
-            else:
-                return "Please respond with 'yes' or 'no'. Would you like to add any extras to your order? (yes/no)"
-
-        # --- Step 6: Capturing Extras Details ---
-        if state.step == 6:
-            state.data["extras"] = message
-            state.step = 7
-            dish = state.data.get("dish", "")
-            quantity = state.data.get("quantity", 1)
-            phone = state.data.get("phone_number", "")
-            address = state.data.get("address", "")
-            shipping_cost = state.data.get("shipping_cost", 0)
-            extras = state.data.get("extras", "")
-            price_per_serving = 1500  # Fixed price per serving
-            total_price = (quantity * price_per_serving) + shipping_cost
-            summary = (f"Order Summary:\nDish: {dish}\nQuantity: {quantity}\n"
-                       f"Phone: {phone}\nAddress: {address}\n"
-                       f"Shipping Cost: N{shipping_cost}\n"
-                       f"Total Price: N{total_price}\n"
-                       f"Extras: {extras}\nConfirm order? (yes/no)")
-            return summary
-
-        # --- Step 7: Order Confirmation & Finalization ---
-        if state.step == 7:
-            if message.lower() in ["yes", "y"]:
-                order_id = f"ORD-{int(time.time())}"
-                state.data["order_id"] = order_id
-                price_per_serving = 1500  # Fixed price per serving
-                quantity = state.data.get("quantity", 1)
-                shipping_cost = state.data.get("shipping_cost", 0)
-                total_price = (quantity * price_per_serving) + shipping_cost
-                state.data["price"] = str(total_price)
-
-                # Save the order asynchronously (including delivery_address, phone, extras, shipping_cost)
-                async def save_order():
-                    async with async_session() as session:
-                        order = Order(
-                            order_id=order_id,
-                            user_id=user_id,
-                            dish=state.data["dish"],
-                            quantity=str(quantity),
-                            price=str(total_price),
-                            status="Pending Payment",
-                            delivery_address=state.data.get("address", ""),
-                            shipping_cost=str(shipping_cost)  # Save shipping cost
-                        )
-                        session.add(order)
-                        await session.commit()
-                asyncio.create_task(save_order())
-
-                # Record the initial tracking update: Order Placed
-                await log_order_tracking(order_id, "Order Placed", "Order placed and awaiting payment.")
-
-                # Notify management of the new order via WhatsApp
-                async def notify_management_order(order_details: dict):
-                    message_body = (
-                        f"New Order Received:\n"
-                        f"Order ID: {order_details['order_id']}\n"
-                        f"Dish: {order_details['dish']}\n"
-                        f"Quantity: {order_details['quantity']}\n"
-                        f"Total Price: {order_details['price']}\n"
-                        f"Phone: {state.data.get('phone_number', '')}\n"
-                        f"Delivery Address: {order_details.get('address', 'Not Provided')}\n"
-                        f"Extras: {state.data.get('extras', 'None')}\n"
-                        f"Status: Pending Payment"
-                    )
-                    await asyncio.to_thread(send_whatsapp_message, MANAGEMENT_WHATSAPP_NUMBER, message_body)
-                order_details = {
-                    "order_id": order_id,
-                    "dish": state.data["dish"],
-                    "quantity": state.data["quantity"],
-                    "price": state.data["price"],
-                    "address": state.data.get("address", "")
-                }
-                asyncio.create_task(notify_management_order(order_details))
-
-                # Update user profile with the new order ID
-                await update_user_profile_with_order(user_id, order_id)
-
-                # Generate payment link
-                email = "customer@example.com"  # Placeholder; retrieve from profile if available
-                payment_data = create_paystack_payment_link(email, total_price * 100, order_id)
-                dish_name = state.data.get("dish", "")
-                state.reset()
-                if user_id in user_state:
-                    del user_state[user_id]
-                if payment_data.get("status"):
-                    payment_link = payment_data["data"]["authorization_url"]
-                    return (f"Thank you for your order of {quantity} serving(s) of {dish_name}! "
-                            f"Your Order ID is {order_id}.\nPlease complete payment here: {payment_link}\n"
-                            "You can track your order status using your Order ID.\n"
-                            "Is there anything else you'd like to order?")
-                else:
-                    return f"Your order has been placed with Order ID {order_id}, but we could not initialize payment. Please try again later."
-            else:
-                state.reset()
-                if user_id in user_state:
-                    del user_state[user_id]
-                return "Order canceled. Let me know if you'd like to try again."
-
-    return ""
-    
+    """
+    Calculate shipping cost based on the user's address.
+    """
+    address_lower = address.lower()
+    for area, cost in TOWN_SHIPPING_COSTS.items():
+        if area in address_lower:
+            return cost
+    return TOWN_SHIPPING_COSTS["default"]  # Default shipping cost for unknown areas
+
+# ... (rest of your imports and setup)
+
+async def update_user_profile(user_id: str, phone_number: str = None, address: str = None) -> UserProfile:
+    """
+    Update or create a user profile with the provided phone number and address.
+    """
+    async with async_session() as session:
+        result = await session.execute(
+            select(UserProfile).where(UserProfile.user_id == user_id)
+        )
+        profile = result.scalars().first()
+        if profile is None:
+            profile = UserProfile(
+                user_id=user_id,
+                phone_number=phone_number,
+                address=address,
+                last_interaction=datetime.utcnow()
+            )
+        else:
+            if phone_number:
+                profile.phone_number = phone_number
+            if address:
+                profile.address = address
+            profile.last_interaction = datetime.utcnow()
+        session.add(profile)
+        await session.commit()
+        return profile
+
+# ... (rest of your imports and setup)
+
+def process_order_flow(user_id: str, message: str) -> str:
+    """
+    Implements an FSM-based order flow with shipping cost calculation.
+    """
+            
+    # Retrieve or initialize conversation state
+    state = user_state.get(user_id)
+    if state and state.is_expired():
+        state.reset()
+        del user_state[user_id]
+        state = None
+
+    # Start a new order flow if the user explicitly types "order" or "menu"
+    if message.lower() in ["order", "menu"]:
+        state = ConversationState()
+        state.flow = "order"
+        state.step = 1
+        state.update_last_active()
+        user_state[user_id] = state
+        if message.lower() == "order":
+            return "Sure! What dish would you like to order?"
+        return ""
+
+    # If no state exists but the message includes "order", start the order flow.
+    if not state and "order" in message.lower():
+        state = ConversationState()
+        state.flow = "order"
+        state.step = 1
+        state.update_last_active()
+        user_state[user_id] = state
+        return "Sure! What dish would you like to order?"
+
+    # --- New Logic: Parse Dish, Quantity, Phone Number, and Address in a Single Message ---
+    if not state or state.flow != "order":
+        # Check if the message contains a dish and quantity
+        dish_candidates = [item["name"] for item in menu_items]
+        found_dish = None
+        for dish in dish_candidates:
+            if dish.lower() in message.lower():
+                found_dish = dish
+                break
+        numbers = re.findall(r'\d+', message)
+        if found_dish and numbers:
+            quantity = int(numbers[0])
+            if quantity <= 0:
+                return "Please enter a valid quantity (e.g., 1, 2, 3)."
+            # Initialize state and skip to step 3 (phone number)
+            state = ConversationState()
+            state.flow = "order"
+            state.step = 3
+            state.data["dish"] = found_dish
+            state.data["quantity"] = quantity
+            state.update_last_active()
+            user_state[user_id] = state
+
+            # Extract phone number and address from the message
+            phone_pattern = r'(\+?\d{10,15})'
+            phone_match = re.search(phone_pattern, message)
+            address = None
+            if phone_match:
+                phone_number = phone_match.group(1)
+                # Assume the address starts after the phone number
+                address_start = phone_match.end()
+                address = message[address_start:].strip()
+                # Remove any leading/trailing commas or spaces
+                address = re.sub(r'^[,\s]+|[,\s]+$', '', address)
+            
+            if phone_match and address:
+                state.data["phone_number"] = phone_number
+                state.data["address"] = address
+                # Save phone number and address to the user's profile
+                asyncio.create_task(update_user_profile(user_id, phone_number, address))
+                # Calculate shipping cost based on the address
+                shipping_cost = calculate_shipping_cost(address)
+                state.data["shipping_cost"] = shipping_cost
+                state.step = 5
+                return (f"Thanks! Your phone number is recorded as: {phone_number}.\n"
+                        f"Your delivery address is: {address}.\n"
+                        f"Your delivery cost is N{shipping_cost}. Would you like to add any extras such as sides or drinks? (yes/no)")
+            elif phone_match:
+                state.data["phone_number"] = phone_match.group(1)
+                # Save phone number to the user's profile
+                asyncio.create_task(update_user_profile(user_id, phone_number))
+                return "Thank you. Please provide your delivery address."
+            else:
+                return "Please provide both your phone number and delivery address. For example: '09162409591, 1, Iyana Isashi, Isashi, Ojo, Lagos'."
+
+    if state and state.flow == "order":
+        state.update_last_active()
+        # --- Step 1: Expecting Dish Selection (and optionally quantity) ---
+        if state.step == 1:
+            dish_candidates = [item["name"] for item in menu_items]
+            found_dish = None
+            for dish in dish_candidates:
+                if dish.lower() in message.lower():
+                    found_dish = dish
+                    break
+            numbers = re.findall(r'\d+', message)
+            if found_dish:
+                state.data["dish"] = found_dish
+                if numbers:
+                    quantity = int(numbers[0])
+                    if quantity <= 0:
+                        return "Please enter a valid quantity (e.g., 1, 2, 3)."
+                    state.data["quantity"] = quantity
+                    state.step = 3  # Move to phone number step
+                    return f"You selected {found_dish} with {quantity} serving(s). Please provide your phone number and delivery address."
+                else:
+                    state.step = 2  # Ask for quantity
+                    return f"You selected {found_dish}. How many servings would you like?"
+            else:
+                return "I couldn't identify the dish. Please type the dish name from our menu."
+
+        # --- Step 2: Asking for Quantity ---
+        if state.step == 2:
+            numbers = re.findall(r'\d+', message)
+            if not numbers:
+                return "Please enter a valid number for the quantity (e.g., 1, 2, 3)."
+            quantity = int(numbers[0])
+            if quantity <= 0:
+                return "Please enter a valid quantity (e.g., 1, 2, 3)."
+            state.data["quantity"] = quantity
+            state.step = 3
+            return f"Got it. {quantity} serving(s) of {state.data.get('dish')}. Please provide your phone number and delivery address."
+
+        # --- Step 3: Parse Phone Number and Address from Message ---
+        if state.step == 3:
+            # Extract phone number using regex
+            phone_pattern = r'(\+?\d{10,15})'
+            phone_match = re.search(phone_pattern, message)
+            
+            # Extract address (assume it follows the phone number, separated by a comma)
+            address = None
+            if phone_match:
+                phone_number = phone_match.group(1)
+                # Assume the address starts after the phone number
+                address_start = phone_match.end()
+                address = message[address_start:].strip()
+                # Remove any leading/trailing commas or spaces
+                address = re.sub(r'^[,\s]+|[,\s]+$', '', address)
+            
+            if phone_match and address:
+                state.data["phone_number"] = phone_number
+                state.data["address"] = address
+                # Save phone number and address to the user's profile
+                asyncio.create_task(update_user_profile(user_id, phone_number, address))
+                # Calculate shipping cost based on the address
+                shipping_cost = calculate_shipping_cost(address)
+                state.data["shipping_cost"] = shipping_cost
+                state.step = 5
+                return (f"Thanks! Your phone number is recorded as: {phone_number}.\n"
+                        f"Your delivery address is: {address}.\n"
+                        f"Your delivery cost is N{shipping_cost}. Would you like to add any extras such as sides or drinks? (yes/no)")
+            elif phone_match:
+                state.data["phone_number"] = phone_match.group(1)
+                # Save phone number to the user's profile
+                asyncio.create_task(update_user_profile(user_id, phone_number))
+                return "Thank you. Please provide your delivery address."
+            else:
+                return "Please provide both your phone number and delivery address. For example: '09162409591, 1, Iyana Isashi, Isashi, Ojo, Lagos'."
+
+        # --- Step 4: Requesting Delivery Address ---
+        if state.step == 4:
+            state.data["address"] = message
+            # Save address to the user's profile
+            asyncio.create_task(update_user_profile(user_id, address=message))
+            # Calculate shipping cost based on the address
+            shipping_cost = calculate_shipping_cost(message)
+            state.data["shipping_cost"] = shipping_cost
+            state.step = 5
+            return (f"Thanks. Your delivery address is recorded as: {message}.\n"
+                    f"Your delivery cost is N{shipping_cost}. Would you like to add any extras such as sides or drinks? (yes/no)")
+
+        # --- Step 5: Asking if They Want Extras ---
+        if state.step == 5:
+            if message.lower() in ["yes", "y"]:
+                state.step = 6
+                return "Please list the extras you would like to add (e.g., drinks, sides, etc.)."
+            elif message.lower() in ["no", "n"]:
+                state.data["extras"] = ""
+                state.step = 7
+                dish = state.data.get("dish", "")
+                quantity = state.data.get("quantity", 1)
+                phone = state.data.get("phone_number", "")
+                address = state.data.get("address", "")
+                shipping_cost = state.data.get("shipping_cost", 0)
+                price_per_serving = 1500  # Fixed price per serving
+                total_price = (quantity * price_per_serving) + shipping_cost
+                summary = (f"Order Summary:\nDish: {dish}\nQuantity: {quantity}\n"
+                           f"Phone: {phone}\nAddress: {address}\n"
+                           f"Shipping Cost: N{shipping_cost}\n"
+                           f"Total Price: N{total_price}\n"
+                           f"Extras: None\nConfirm order? (yes/no)")
+                return summary
+            else:
+                return "Please respond with 'yes' or 'no'. Would you like to add any extras to your order? (yes/no)"
+
+        # --- Step 6: Capturing Extras Details ---
+        if state.step == 6:
+            state.data["extras"] = message
+            state.step = 7
+            dish = state.data.get("dish", "")
+            quantity = state.data.get("quantity", 1)
+            phone = state.data.get("phone_number", "")
+            address = state.data.get("address", "")
+            shipping_cost = state.data.get("shipping_cost", 0)
+            extras = state.data.get("extras", "")
+            price_per_serving = 1500  # Fixed price per serving
+            total_price = (quantity * price_per_serving) + shipping_cost
+            summary = (f"Order Summary:\nDish: {dish}\nQuantity: {quantity}\n"
+                       f"Phone: {phone}\nAddress: {address}\n"
+                       f"Shipping Cost: N{shipping_cost}\n"
+                       f"Total Price: N{total_price}\n"
+                       f"Extras: {extras}\nConfirm order? (yes/no)")
+            return summary
+
+        # --- Step 7: Order Confirmation & Finalization ---
+        if state.step == 7:
+            if message.lower() in ["yes", "y"]:
+                order_id = f"ORD-{int(time.time())}"
+                state.data["order_id"] = order_id
+                price_per_serving = 1500  # Fixed price per serving
+                quantity = state.data.get("quantity", 1)
+                shipping_cost = state.data.get("shipping_cost", 0)
+                total_price = (quantity * price_per_serving) + shipping_cost
+                state.data["price"] = str(total_price)
+
+                # Save the order asynchronously (including delivery_address, phone, extras, shipping_cost)
+                async def save_order():
+                    async with async_session() as session:
+                        order = Order(
+                            order_id=order_id,
+                            user_id=user_id,
+                            dish=state.data["dish"],
+                            quantity=str(quantity),
+                            price=str(total_price),
+                            status="Pending Payment",
+                            delivery_address=state.data.get("address", ""),
+                            shipping_cost=str(shipping_cost)  # Save shipping cost
+                        )
+                        session.add(order)
+                        await session.commit()
+                asyncio.create_task(save_order())
+
+                # Record the initial tracking update: Order Placed
+                asyncio.create_task(log_order_tracking(order_id, "Order Placed", "Order placed and awaiting payment."))
+
+                # Notify management of the new order via WhatsApp
+                async def notify_management_order(order_details: dict):
+                    message_body = (
+                        f"New Order Received:\n"
+                        f"Order ID: {order_details['order_id']}\n"
+                        f"Dish: {order_details['dish']}\n"
+                        f"Quantity: {order_details['quantity']}\n"
+                        f"Total Price: {order_details['price']}\n"
+                        f"Phone: {state.data.get('phone_number', '')}\n"
+                        f"Delivery Address: {order_details.get('address', 'Not Provided')}\n"
+                        f"Extras: {state.data.get('extras', 'None')}\n"
+                        f"Status: Pending Payment"
+                    )
+                    await asyncio.to_thread(send_whatsapp_message, MANAGEMENT_WHATSAPP_NUMBER, message_body)
+                order_details = {
+                    "order_id": order_id,
+                    "dish": state.data["dish"],
+                    "quantity": state.data["quantity"],
+                    "price": state.data["price"],
+                    "address": state.data.get("address", "")
+                }
+                asyncio.create_task(notify_management_order(order_details))
+
+                email = "customer@example.com"  # Placeholder; retrieve from profile if available
+                payment_data = create_paystack_payment_link(email, total_price * 100, order_id)
+                dish_name = state.data.get("dish", "")
+                state.reset()
+                if user_id in user_state:
+                    del user_state[user_id]
+                if payment_data.get("status"):
+                    payment_link = payment_data["data"]["authorization_url"]
+                    return (f"Thank you for your order of {quantity} serving(s) of {dish_name}! "
+                            f"Your Order ID is {order_id}.\nPlease complete payment here: {payment_link}\n"
+                            "You can track your order status using your Order ID.\n"
+                            "Is there anything else you'd like to order?")
+                else:
+                    return f"Your order has been placed with Order ID {order_id}, but we could not initialize payment. Please try again later."
+            else:
+                state.reset()
+                if user_id in user_state:
+                    del user_state[user_id]
+                return "Order canceled. Let me know if you'd like to try again."
+
+    return ""
+    
 # --- User Profile Functions ---
 async def get_or_create_user_profile(user_id: str, phone_number: str = None) -> UserProfile:
-    async with async_session() as session:
-        result = await session.execute(
-            select(UserProfile).where(UserProfile.user_id == user_id)
-        )
-        profile = result.scalars().first()
-        if profile is None:
-            profile = UserProfile(
-                user_id=user_id,
-                phone_number=phone_number,
-                last_interaction=datetime.utcnow()
-            )
-            session.add(profile)
-            await session.commit()
-        return profile
+    async with async_session() as session:
+        result = await session.execute(
+            select(UserProfile).where(UserProfile.user_id == user_id)
+        )
+        profile = result.scalars().first()
+        if profile is None:
+            profile = UserProfile(
+                user_id=user_id,
+                phone_number=phone_number,
+                last_interaction=datetime.utcnow()
+            )
+            session.add(profile)
+            await session.commit()
+        return profile
 
 async def update_user_last_interaction(user_id: str):
-    async with async_session() as session:
-        result = await session.execute(
-            select(UserProfile).where(UserProfile.user_id == user_id)
-        )
-        profile = result.scalars().first()
-        if profile:
-            profile.last_interaction = datetime.utcnow()
-            await session.commit()
+    async with async_session() as session:
+        result = await session.execute(
+            select(UserProfile).where(UserProfile.user_id == user_id)
+        )
+        profile = result.scalars().first()
+        if profile:
+            profile.last_interaction = datetime.utcnow()
+            await session.commit()
 
 # --- Proactive Engagement: Warm Greetings ---
 async def send_proactive_greeting(user_id: str):
-    greeting = "Hi again! We miss you. Would you like to see our new menu items or get personalized recommendations?"
-    await log_chat_to_db(user_id, "outbound", greeting)
-    return greeting
+    greeting = "Hi again! We miss you. Would you like to see our new menu items or get personalized recommendations?"
+    await log_chat_to_db(user_id, "outbound", greeting)
+    return greeting
 
 # --- FastAPI Setup & Endpoints ---
 app = FastAPI()
 
 @app.on_event("startup")
 async def on_startup():
-    await init_db()
+    await init_db()
 
-@app.post("/chatbot")
 @app.post("/chatbot")
 async def chatbot_response(request: Request, background_tasks: BackgroundTasks):
-    data = await request.json()
-    user_id = data.get("user_id")
-    user_message = data.get("message", "").strip()
-
-    # Initialize conversation context for the user if not present
-    if user_id not in conversation_context:
-        conversation_context[user_id] = []
-
-    # Append the user's message to the conversation context
-    conversation_context[user_id].append({
-        "timestamp": datetime.utcnow().isoformat(),
-        "role": "user",
-        "message": user_message
-    })
-
-    # Log the user's message
-    background_tasks.add_task(log_chat_to_db, user_id, "inbound", user_message)
-
-    # Analyze sentiment
-    sentiment_score = analyze_sentiment(user_message)
-    background_tasks.add_task(log_sentiment, user_id, user_message, sentiment_score)
-    sentiment_modifier = "Great to hear from you! " if sentiment_score > 0.3 else ""
-
-    # --- Handle Menu Selection ---
-    if user_message.strip() == "1" or "menu" in user_message.lower():
-        if user_id in user_state:
-            del user_state[user_id]  # Clear any existing state
-        menu_with_images = []
-        for index, item in enumerate(menu_items, start=1):
-            image_url = google_image_scrape(item["name"])
-            menu_with_images.append({
-                "number": index,
-                "name": item["name"],
-                "description": item["description"],
-                "price": item["price"],
-                "image_url": image_url
-            })
-        response_payload = {
-            "response": sentiment_modifier + "Here’s our delicious menu:",
-            "menu": menu_with_images,
-            "follow_up": (
-                "To order, type the *number* or *name* of the dish you'd like. "
-                "For example, type '1' or 'Jollof Rice' to order Jollof Rice.\n\n"
-                "You can also ask for nutritional facts by typing, for example, 'Nutritional facts for Jollof Rice'."
-            )
-        }
-        background_tasks.add_task(log_chat_to_db, user_id, "outbound", str(response_payload))
-        conversation_context[user_id].append({
-            "timestamp": datetime.utcnow().isoformat(),
-            "role": "bot",
-            "message": response_payload["response"]
-        })
-        return JSONResponse(content=response_payload)
-
-    # --- Order Flow Handling ---
-    if is_order_intent(user_message) or (user_id in user_state and user_state[user_id].flow == "order"):
-        order_response = await process_order_flow(user_id, user_message)
-        if order_response:
-            background_tasks.add_task(log_chat_to_db, user_id, "outbound", order_response)
-            conversation_context[user_id].append({
-                "timestamp": datetime.utcnow().isoformat(),
-                "role": "bot",
-                "message": order_response
-            })
-            return JSONResponse(content={"response": sentiment_modifier + order_response})
-
-    # --- Fallback to NVIDIA LLM ---
-    recent_context = conversation_context.get(user_id, [])[-5:]  # Get the last 5 messages
-    context_str = "\n".join([f"{entry['role'].capitalize()}: {entry['message']}" for entry in recent_context])
-    prompt = f"Conversation context:\n{context_str}\nUser query: {user_message}\nGenerate a helpful, personalized response for a restaurant chatbot."
-    response_stream = stream_text_completion(prompt)  # Pass only the prompt
-    fallback_response = "".join([chunk for chunk in response_stream])
-
-    # Log the bot's response
-    background_tasks.add_task(log_chat_to_db, user_id, "outbound", fallback_response)
-    conversation_context[user_id].append({
-        "timestamp": datetime.utcnow().isoformat(),
-        "role": "bot",
-        "message": fallback_response
-    })
-
-    return JSONResponse(content={"response": sentiment_modifier + fallback_response})
+    data = await request.json()
+    user_id = data.get("user_id")
+    phone_number = data.get("phone_number")
+    user_message = data.get("message", "").strip()
+    is_image = data.get("is_image", False)
+    image_b64 = data.get("image_base64", None)
+
+    if not user_id:
+        raise HTTPException(status_code=400, detail="Missing user_id in payload.")
+
+    # Initialize conversation context for the user if not present.
+    if user_id not in conversation_context:
+        conversation_context[user_id] = []
+    # Append the inbound message to the conversation context.
+    conversation_context[user_id].append({
+        "timestamp": datetime.utcnow().isoformat(),
+        "role": "user",
+        "message": user_message
+    })
+
+    background_tasks.add_task(log_chat_to_db, user_id, "inbound", user_message)
+    await update_user_last_interaction(user_id)
+    await get_or_create_user_profile(user_id, phone_number)
+
+    # Handle image queries
+    if is_image and image_b64:
+        if len(image_b64) >= 180_000:
+            raise HTTPException(status_code=400, detail="Image too large.")
+        return StreamingResponse(stream_image_completion(image_b64), media_type="text/plain")
+
+    sentiment_score = analyze_sentiment(user_message)
+    background_tasks.add_task(log_sentiment, user_id, user_message, sentiment_score)
+    sentiment_modifier = ""
+    if sentiment_score < -0.3:
+        sentiment_modifier = "I'm sorry if you're having a tough time. "
+    elif sentiment_score > 0.3:
+        sentiment_modifier = "Great to hear from you! "
+
+          # --- Order Tracking Handling ---
+    order_id_match = re.search(r"ord-\d+", user_message.lower())
+    if order_id_match:
+        order_id = order_id_match.group(0)
+        try:
+            # Call the /track_order endpoint
+            tracking_response = await track_order(order_id)
+            return JSONResponse(content={"response": tracking_response})
+        except HTTPException as e:
+            return JSONResponse(content={"response": f" {e.detail}"})
+
+    # --- Order Flow Handling ---
+    order_response = process_order_flow(user_id, user_message)
+    if order_response:
+        background_tasks.add_task(log_chat_to_db, user_id, "outbound", order_response)
+        conversation_context[user_id].append({
+            "timestamp": datetime.utcnow().isoformat(),
+            "role": "bot",
+            "message": order_response
+        })
+        return JSONResponse(content={"response": sentiment_modifier + order_response})
+
+    # --- Menu Display ---
+    if "menu" in user_message.lower():
+        if user_id in user_state:
+            del user_state[user_id]
+        menu_with_images = []
+        for index, item in enumerate(menu_items, start=1):
+            image_url = google_image_scrape(item["name"])
+            menu_with_images.append({
+                "number": index,
+                "name": item["name"],
+                "description": item["description"],
+                "price": item["price"],
+                "image_url": image_url
+            })
+        response_payload = {
+            "response": sentiment_modifier + "Here’s our delicious menu:",
+            "menu": menu_with_images,
+            "follow_up": (
+                "To order, type the *number* or *name* of the dish you'd like. "
+                "For example, type '1' or 'Jollof Rice' to order Jollof Rice.\n\n"
+                "You can also ask for nutritional facts by typing, for example, 'Nutritional facts for Jollof Rice'."
+            )
+        }
+        background_tasks.add_task(log_chat_to_db, user_id, "outbound", str(response_payload))
+        conversation_context[user_id].append({
+            "timestamp": datetime.utcnow().isoformat(),
+            "role": "bot",
+            "message": response_payload["response"]
+        })
+        return JSONResponse(content=response_payload)
+
+    # --- Dish Selection via Menu ---
+    if any(item["name"].lower() in user_message.lower() for item in menu_items) or \
+       any(str(index) == user_message.strip() for index, item in enumerate(menu_items, start=1)):
+        selected_dish = None
+        if user_message.strip().isdigit():
+            dish_number = int(user_message.strip())
+            if 1 <= dish_number <= len(menu_items):
+                selected_dish = menu_items[dish_number - 1]["name"]
+        else:
+            for item in menu_items:
+                if item["name"].lower() in user_message.lower():
+                    selected_dish = item["name"]
+                    break
+        if selected_dish:
+            state = ConversationState()
+            state.flow = "order"
+            # Set step to 2 since the dish is already selected
+            state.step = 2  
+            state.data["dish"] = selected_dish
+            state.update_last_active()
+            user_state[user_id] = state
+            response_text = f"You selected {selected_dish}. How many servings would you like?"
+            background_tasks.add_task(log_chat_to_db, user_id, "outbound", response_text)
+            conversation_context[user_id].append({
+                "timestamp": datetime.utcnow().isoformat(),
+                "role": "bot",
+                "message": response_text
+            })
+            return JSONResponse(content={"response": sentiment_modifier + response_text})
+        else:
+            response_text = "Sorry, I couldn't find that dish in the menu. Please try again."
+            background_tasks.add_task(log_chat_to_db, user_id, "outbound", response_text)
+            conversation_context[user_id].append({
+                "timestamp": datetime.utcnow().isoformat(),
+                "role": "bot",
+                "message": response_text
+            })
+            return JSONResponse(content={"response": sentiment_modifier + response_text})
+
+    # --- Nutritional Facts ---
+    if "nutritional facts for" in user_message.lower():
+        dish_name = user_message.lower().replace("nutritional facts for", "").strip().title()
+        dish = next((item for item in menu_items if item["name"].lower() == dish_name.lower()), None)
+        if dish:
+            response_text = f"Nutritional facts for {dish['name']}:\n{dish['nutrition']}"
+        else:
+            response_text = f"Sorry, I couldn't find nutritional facts for {dish_name}."
+        background_tasks.add_task(log_chat_to_db, user_id, "outbound", response_text)
+        conversation_context[user_id].append({
+            "timestamp": datetime.utcnow().isoformat(),
+            "role": "bot",
+            "message": response_text
+        })
+        return JSONResponse(content={"response": sentiment_modifier + response_text})
+
+    # --- Fallback: LLM Response Streaming with Conversation Context ---
+    recent_context = conversation_context.get(user_id, [])[-5:]
+    context_str = "\n".join([f"{entry['role'].capitalize()}: {entry['message']}" for entry in recent_context])
+    prompt = f"Conversation context:\n{context_str}\nUser query: {user_message}\nGenerate a helpful, personalized response for a restaurant chatbot."
+    def stream_response():
+        for chunk in stream_text_completion(prompt):
+            yield chunk
+    fallback_log = f"LLM fallback response for prompt: {prompt}"
+    background_tasks.add_task(log_chat_to_db, user_id, "outbound", fallback_log)
+    return StreamingResponse(stream_response(), media_type="text/plain")
 
 # --- Other Endpoints (Chat History, Order Details, User Profile, Analytics, Voice, Payment Callback) ---
 @app.get("/chat_history/{user_id}")
 async def get_chat_history(user_id: str):
-    async with async_session() as session:
-        result = await session.execute(
-            ChatHistory.__table__.select().where(ChatHistory.user_id == user_id)
-        )
-        history = result.fetchall()
-        return [dict(row) for row in history]
+    async with async_session() as session:
+        result = await session.execute(
+            ChatHistory.__table__.select().where(ChatHistory.user_id == user_id)
+        )
+        history = result.fetchall()
+        return [dict(row) for row in history]
 
 @app.get("/order/{order_id}")
 async def get_order(order_id: str):
-    async with async_session() as session:
-        result = await session.execute(
-            Order.__table__.select().where(Order.order_id == order_id)
-        )
-        order = result.fetchone()
-        if order:
-            return dict(order)
-        else:
-            raise HTTPException(status_code=404, detail="Order not found.")
+    async with async_session() as session:
+        result = await session.execute(
+            Order.__table__.select().where(Order.order_id == order_id)
+        )
+        order = result.fetchone()
+        if order:
+            return dict(order)
+        else:
+            raise HTTPException(status_code=404, detail="Order not found.")
 
 @app.get("/user_profile/{user_id}")
 async def get_user_profile(user_id: str):
-    profile = await get_or_create_user_profile(user_id)
-    return {
-        "user_id": profile.user_id,
-        "phone_number": profile.phone_number,
-        "name": profile.name,
-        "email": profile.email,
-        "preferences": profile.preferences,
-        "last_interaction": profile.last_interaction.isoformat(),
-        "order_ids": profile.order_ids
-    }
+    profile = await get_or_create_user_profile(user_id)
+    return {
+        "user_id": profile.user_id,
+        "phone_number": profile.phone_number,
+        "name": profile.name,
+        "email": profile.email,
+        "preferences": profile.preferences,
+        "last_interaction": profile.last_interaction.isoformat()
+    }
 
 @app.get("/analytics")
 async def get_analytics():
-    async with async_session() as session:
-        msg_result = await session.execute(ChatHistory.__table__.count())
-        total_messages = msg_result.scalar() or 0
-        order_result = await session.execute(Order.__table__.count())
-        total_orders = order_result.scalar() or 0
-        sentiment_result = await session.execute("SELECT AVG(sentiment_score) FROM sentiment_logs")
-        avg_sentiment = sentiment_result.scalar() or 0
-    return {
-        "total_messages": total_messages,
-        "total_orders": total_orders,
-        "average_sentiment": avg_sentiment
-    }
+    async with async_session() as session:
+        msg_result = await session.execute(ChatHistory.__table__.count())
+        total_messages = msg_result.scalar() or 0
+        order_result = await session.execute(Order.__table__.count())
+        total_orders = order_result.scalar() or 0
+        sentiment_result = await session.execute("SELECT AVG(sentiment_score) FROM sentiment_logs")
+        avg_sentiment = sentiment_result.scalar() or 0
+    return {
+        "total_messages": total_messages,
+        "total_orders": total_orders,
+        "average_sentiment": avg_sentiment
+    }
 
 # Load the Hugging Face API token from environment variables
 HUGGING_FACE_API_TOKEN = os.getenv("HUGGING_FACE_API_TOKEN")
 if not HUGGING_FACE_API_TOKEN:
-    raise ValueError("Hugging Face API token not found in environment variables.")
+    raise ValueError("Hugging Face API token not found in environment variables.")
 
 # Hugging Face Whisper API configuration
 WHISPER_API_URL = "https://router.huggingface.co/fal-ai"
 WHISPER_API_HEADERS = {"Authorization": f"Bearer {HUGGING_FACE_API_TOKEN}"}
 
 class TranscriptionResponse(BaseModel):
-    transcription: str
+    transcription: str
 
 @app.post("/voice", response_model=TranscriptionResponse)
 async def process_voice(file: UploadFile = File(...)):
-    """
-    Endpoint to process voice notes and transcribe them using OpenAI Whisper via Hugging Face.
-    """
-    try:
-        # Read the uploaded file
-        contents = await file.read()
+    """
+    Endpoint to process voice notes and transcribe them using OpenAI Whisper via Hugging Face.
+    """
+    try:
+        # Read the uploaded file
+        contents = await file.read()
 
-        # Save the file temporarily (optional, depending on how the API expects the input)
-        temp_file_path = f"temp_{file.filename}"
-        with open(temp_file_path, "wb") as temp_file:
-            temp_file.write(contents)
+        # Save the file temporarily (optional, depending on how the API expects the input)
+        temp_file_path = f"temp_{file.filename}"
+        with open(temp_file_path, "wb") as temp_file:
+            temp_file.write(contents)
 
-        # Call the Hugging Face Whisper API
-        with open(temp_file_path, "rb") as audio_file:
-            response = requests.post(
-                WHISPER_API_URL,
-                headers=WHISPER_API_HEADERS,
-                files={"file": audio_file}
-            )
+        # Call the Hugging Face Whisper API
+        with open(temp_file_path, "rb") as audio_file:
+            response = requests.post(
+                WHISPER_API_URL,
+                headers=WHISPER_API_HEADERS,
+                files={"file": audio_file}
+            )
 
-        # Clean up the temporary file
-        os.remove(temp_file_path)
+        # Clean up the temporary file
+        os.remove(temp_file_path)
 
-        # Check if the API call was successful
-        if response.status_code != 200:
-            raise HTTPException(status_code=response.status_code, detail="Failed to transcribe audio.")
+        # Check if the API call was successful
+        if response.status_code != 200:
+            raise HTTPException(status_code=response.status_code, detail="Failed to transcribe audio.")
 
-        # Extract the transcription from the API response
-        transcription = response.json().get("text", "")
+        # Extract the transcription from the API response
+        transcription = response.json().get("text", "")
 
-        return {"transcription": transcription}
+        return {"transcription": transcription}
 
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
 
 # --- Payment Callback Endpoint with Payment Tracking and Redirection ---
 @app.api_route("/payment_callback", methods=["GET", "POST"])
-async def payment_callback(request: Request):
-    # GET: User redirection after payment
-    if request.method == "GET":
-        params = request.query_params
-        order_id = params.get("reference")
-        status = params.get("status", "Paid")
-        if not order_id:
-            raise HTTPException(status_code=400, detail="Missing order reference in callback.")
-        async with async_session() as session:
-            result = await session.execute(
-                Order.__table__.select().where(Order.order_id == order_id)
-            )
-            order = result.scalar_one_or_none()
-            if order:
-                order.status = status
-                await session.commit()
-            else:
-                raise HTTPException(status_code=404, detail="Order not found.")
-        # Record payment confirmation tracking update
-        await log_order_tracking(order_id, "Payment Confirmed", f"Payment status updated to {status}.")
-        # Notify management via WhatsApp about the payment update
-        await asyncio.to_thread(send_whatsapp_message, MANAGEMENT_WHATSAPP_NUMBER,
-            f"Payment Update:\nOrder ID: {order_id} is now {status}."
-        )
-        # Redirect user back to the chat interface (adjust URL as needed)
-        redirect_url = f"https://wa.link/am87s2"
-        return RedirectResponse(url=redirect_url)
-    # POST: Server-to-server callback from Paystack
-    else:
-        data = await request.json()
-        order_id = data.get("reference")
-        new_status = data.get("status", "Paid")
-        if not order_id:
-            raise HTTPException(status_code=400, detail="Missing order reference in callback.")
-        async with async_session() as session:
-            result = await session.execute(
-                Order.__table__.select().where(Order.order_id == order_id)
-            )
-            order = result.scalar_one_or_none()
-            if order:
-                order.status = new_status
-                await session.commit()
-                await log_order_tracking(order_id, "Payment Confirmed", f"Payment status updated to {new_status}.")
-                await asyncio.to_thread(send_whatsapp_message, MANAGEMENT_WHATSAPP_NUMBER,
-                    f"Payment Update:\nOrder ID: {order_id} is now {new_status}."
-                )
-                return JSONResponse(content={"message": "Order updated successfully."})
-            else:
-                raise HTTPException(status_code=404, detail="Order not found.")
 
+async def payment_callback(request: Request):
+    # GET: User redirection after payment
+    if request.method == "GET":
+        params = request.query_params
+        order_id = params.get("reference")
+        status = params.get("status", "Paid")
+        if not order_id:
+            raise HTTPException(status_code=400, detail="Missing order reference in callback.")
+        async with async_session() as session:
+            result = await session.execute(
+                Order.__table__.select().where(Order.order_id == order_id)
+            )
+            order = result.scalar_one_or_none()
+            if order:
+                order.status = status
+                await session.commit()
+            else:
+                raise HTTPException(status_code=404, detail="Order not found.")
+        # Record payment confirmation tracking update
+        await log_order_tracking(order_id, "Payment Confirmed", f"Payment status updated to {status}.")
+        # Notify management via WhatsApp about the payment update
+        await asyncio.to_thread(send_whatsapp_message, MANAGEMENT_WHATSAPP_NUMBER,
+            f"Payment Update:\nOrder ID: {order_id} is now {status}."
+        )
+        # Redirect user back to the chat interface (adjust URL as needed)
+        redirect_url = f"https://wa.link/am87s2"
+        return RedirectResponse(url=redirect_url)
+    # POST: Server-to-server callback from Paystack
+    else:
+        data = await request.json()
+        order_id = data.get("reference")
+        new_status = data.get("status", "Paid")
+        if not order_id:
+            raise HTTPException(status_code=400, detail="Missing order reference in callback.")
+        async with async_session() as session:
+            result = await session.execute(
+                Order.__table__.select().where(Order.order_id == order_id)
+            )
+            order = result.scalar_one_or_none()
+            if order:
+                order.status = new_status
+                await session.commit()
+                await log_order_tracking(order_id, "Payment Confirmed", f"Payment status updated to {new_status}.")
+                await asyncio.to_thread(send_whatsapp_message, MANAGEMENT_WHATSAPP_NUMBER,
+                    f"Payment Update:\nOrder ID: {order_id} is now {new_status}."
+                )
+                return JSONResponse(content={"message": "Order updated successfully."})
+            else:
+                raise HTTPException(status_code=404, detail="Order not found.")
+                
 @app.get("/track_order/{order_id}")
 async def track_order(order_id: str):
-    """
-    Fetch order tracking details for a given order ID.
-    """
-    async with async_session() as session:
-        result = await session.execute(
-            select(OrderTracking)
-            .where(OrderTracking.order_id == order_id)
-            .order_by(OrderTracking.timestamp)
-        )
-        tracking_updates = result.scalars().all()
-        if tracking_updates:
-            response = []
-            for update in tracking_updates:
-                response.append({
-                    "status": update.status,
-                    "message": update.message,
-                    "timestamp": update.timestamp.isoformat(),
-                })
-            return JSONResponse(content=response)
-        else:
-            raise HTTPException(status_code=404, detail="No tracking information found for this order.")
+    """
+    Fetch order tracking details for a given order ID.
+    """
+    async with async_session() as session:
+        result = await session.execute(
+            select(OrderTracking)
+            .where(OrderTracking.order_id == order_id)
+            .order_by(OrderTracking.timestamp)
+        )
+        tracking_updates = result.scalars().all()
+        if tracking_updates:
+            response = []
+            for update in tracking_updates:
+                response.append({
+                    "status": update.status,
+                    "message": update.message,
+                    "timestamp": update.timestamp.isoformat(),
+                })
+            return JSONResponse(content=response)
+        else:
+            raise HTTPException(status_code=404, detail="No tracking information found for this order.")
 
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)