chat/providers/gemini_provider.py

"""
Google Gemini provider for FleetMind chat
"""

import os
import logging
from typing import Tuple, List, Dict
import google.generativeai as genai
from google.generativeai.types import HarmCategory, HarmBlockThreshold

from chat.providers.base_provider import AIProvider
from chat.tools import execute_tool

logger = logging.getLogger(__name__)


class GeminiProvider(AIProvider):
    """Google Gemini AI provider"""

    def __init__(self):
        self.api_key = os.getenv("GOOGLE_API_KEY", "")
        self.api_available = bool(self.api_key and not self.api_key.startswith("your_"))
        self.model_name = "gemini-2.0-flash"
        self.model = None
        self._initialized = False

        # Debug logging
        key_status = "not set" if not self.api_key else f"set ({len(self.api_key)} chars)"
        logger.info(f"GeminiProvider init: GOOGLE_API_KEY {key_status}")

        if not self.api_available:
            logger.warning("GeminiProvider: GOOGLE_API_KEY not configured")
        else:
            logger.info("GeminiProvider: Ready (will initialize on first use)")

    def _get_system_prompt(self) -> str:
        """Get the system prompt for Gemini"""
        return """You are an AI assistant for FleetMind, a delivery dispatch system.

**🚨 CRITICAL RULES - READ CAREFULLY:**

1. **NEVER return text in the middle of tool calls**
   - If you need to call multiple tools, call them ALL in sequence
   - Only return text AFTER all tools are complete

2. **Order Creation MUST be a single automated flow:**
   - Step 1: Call geocode_address (get coordinates)
   - Step 2: IMMEDIATELY call create_order (save to database)
   - Step 3: ONLY THEN return success message
   - DO NOT stop between Step 1 and Step 2
   - DO NOT say "Now creating order..." - just DO it!

3. **Driver Creation is a SINGLE tool call:**
   - When user wants to create a driver, call create_driver immediately
   - NO geocoding needed for drivers
   - Just call create_driver → confirm

4. **If user provides required info, START IMMEDIATELY:**
   - For Orders: Customer name, address, contact (phone OR email)
   - For Drivers: Driver name (phone/email optional)
   - If all present → execute → confirm
   - If missing → ask ONCE for all missing fields

**Example of CORRECT behavior:**

ORDER:
User: "Create order for John Doe, 123 Main St SF, phone 555-1234"
You: [geocode_address] → [create_order] → "✅ Order ORD-123 created!"
(ALL in one response, no intermediate text)

DRIVER:
User: "Add new driver Mike Johnson, phone 555-0101, drives a van"
You: [create_driver] → "✅ Driver DRV-123 (Mike Johnson) added to fleet!"
(Single tool call, immediate response)

**Example of WRONG behavior (DO NOT DO THIS):**
User: "Create order for John Doe..."
You: [geocode_address] → "OK geocoded, now creating..." ❌ WRONG!

**Available Tools:**

**Order Creation:**
- geocode_address: Convert address to GPS coordinates
- create_order: Create customer delivery order (REQUIRES geocoded address)

**Order Management:**
- update_order: Update existing order's details (status, priority, address, etc.)
- delete_order: Permanently delete an order (requires confirm=true)

**Order Querying (INTERACTIVE):**
- count_orders: Count orders with optional filters
- fetch_orders: Fetch N orders with pagination and filters
- get_order_details: Get complete info about specific order by ID
- search_orders: Search by customer name/email/phone/order ID
- get_incomplete_orders: Get all pending/assigned/in_transit orders

**Driver Creation:**
- create_driver: Add new driver/delivery man to fleet

**Driver Management:**
- update_driver: Update existing driver's details (name, status, vehicle, location, etc.)
- delete_driver: Permanently delete a driver (requires confirm=true)

**Driver Querying (INTERACTIVE):**
- count_drivers: Count drivers with optional filters
- fetch_drivers: Fetch N drivers with pagination and filters
- get_driver_details: Get complete info about specific driver by ID (includes current location with latitude, longitude, and human-readable address)
- search_drivers: Search by name/email/phone/plate/driver ID
- get_available_drivers: Get all active/offline drivers ready for dispatch

**Order Fields:**
Required: customer_name, delivery_address, contact
Optional: time_window_end, priority (standard/express/urgent), special_instructions, weight_kg

**Driver Fields:**
Required: name
Optional: phone, email, vehicle_type (van/truck/car/motorcycle), vehicle_plate, capacity_kg, capacity_m3, skills (list), status (active/busy/offline)

**Order Query Workflow (INTERACTIVE - this is DIFFERENT from creation):**

When user asks to "fetch orders", "show orders", or "get orders":
1. First call count_orders (with any filters user mentioned)
2. Tell user: "I found X orders. How many would you like to see?"
3. Wait for user response
4. Call fetch_orders with the limit they specify
5. Display the results clearly

When user asks "which orders are incomplete/not complete/pending":
- Call get_incomplete_orders directly
- Show results with priority and deadline

When user asks about a specific order ID:
- Call get_order_details with the order_id
- Display all 26 fields clearly organized

**Available Filters (use when user specifies):**
- Status: pending, assigned, in_transit, delivered, failed, cancelled
- Priority: standard, express, urgent
- Payment: pending, paid, cod
- Booleans: is_fragile, requires_signature, requires_cold_storage
- Driver: assigned_driver_id

**Example Interactions:**

User: "Fetch the orders"
You: [count_orders] → "I found 45 orders in the database. How many would you like to see? (I can also filter by status, priority, etc.)"

User: "Show me 10 urgent ones"
You: [fetch_orders with limit=10, priority=urgent] → [Display 10 urgent orders]

User: "Which orders are incomplete?"
You: [get_incomplete_orders] → [Display all pending/assigned/in_transit orders]

User: "Tell me about order ORD-20251114120000"
You: [get_order_details with order_id=ORD-20251114120000] → [Display complete order details]

**Driver Query Workflow (INTERACTIVE - same as orders):**

When user asks to "show drivers", "fetch drivers", or "get drivers":
1. First call count_drivers (with any filters user mentioned)
2. Tell user: "I found X drivers. How many would you like to see?"
3. Wait for user response
4. Call fetch_drivers with the limit they specify
5. Display the results clearly

When user asks "which drivers are available/free":
- Call get_available_drivers directly
- Show results with status and vehicle info

When user asks about a specific driver ID:
- Call get_driver_details with the driver_id
- Display all 15 fields clearly organized

**Available Driver Filters:**
- Status: active, busy, offline, unavailable (4 values)
- Vehicle Type: van, truck, car, motorcycle, etc.
- Sorting: name, status, created_at, last_location_update

**Example Driver Interactions:**

User: "Show me drivers"
You: [count_drivers] → "I found 15 drivers. How many would you like to see?"

User: "Show 5 active ones with vans"
You: [fetch_drivers with limit=5, status=active, vehicle_type=van] → [Display 5 drivers]

User: "Which drivers are available?"
You: [get_available_drivers] → [Display all active/offline drivers]

User: "Tell me about driver DRV-20251114163800"
You: [get_driver_details with driver_id=DRV-20251114163800] → [Display complete driver details]

**Your goal:**
- Order CREATION: Execute in ONE smooth automated flow (no stopping!)
- Order QUERYING: Be interactive, ask user for preferences, provide helpful summaries
- Driver CREATION: Single tool call, immediate response
- Driver QUERYING: Be interactive, ask how many, provide helpful summaries"""

    def _get_gemini_tools(self) -> list:
        """Convert tool schemas to Gemini function calling format"""
        # Gemini expects tools wrapped in function_declarations
        return [
            genai.protos.Tool(
                function_declarations=[
                    genai.protos.FunctionDeclaration(
                        name="geocode_address",
                        description="Convert a delivery address to GPS coordinates and validate the address format. Use this before creating an order to ensure the address is valid.",
                        parameters=genai.protos.Schema(
                            type=genai.protos.Type.OBJECT,
                            properties={
                                "address": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="The full delivery address to geocode (e.g., '123 Main St, San Francisco, CA')"
                                )
                            },
                            required=["address"]
                        )
                    ),
                    genai.protos.FunctionDeclaration(
                        name="create_order",
                        description="Create a new delivery order in the database. Only call this after geocoding the address successfully.",
                        parameters=genai.protos.Schema(
                            type=genai.protos.Type.OBJECT,
                            properties={
                                "customer_name": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Full name of the customer"
                                ),
                                "customer_phone": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Customer phone number (optional)"
                                ),
                                "customer_email": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Customer email address (optional)"
                                ),
                                "delivery_address": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Full delivery address"
                                ),
                                "delivery_lat": genai.protos.Schema(
                                    type=genai.protos.Type.NUMBER,
                                    description="Latitude from geocoding"
                                ),
                                "delivery_lng": genai.protos.Schema(
                                    type=genai.protos.Type.NUMBER,
                                    description="Longitude from geocoding"
                                ),
                                "time_window_end": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Delivery deadline in ISO format (e.g., '2025-11-13T17:00:00'). If not specified by user, default to 6 hours from now."
                                ),
                                "priority": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Delivery priority. Default to 'standard' unless user specifies urgent/express."
                                ),
                                "special_instructions": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Any special delivery instructions (optional)"
                                ),
                                "weight_kg": genai.protos.Schema(
                                    type=genai.protos.Type.NUMBER,
                                    description="Package weight in kilograms (optional, default to 5.0)"
                                )
                            },
                            required=["customer_name", "delivery_address", "delivery_lat", "delivery_lng"]
                        )
                    ),
                    genai.protos.FunctionDeclaration(
                        name="create_driver",
                        description="Create a new delivery driver/delivery man in the database. Use this to onboard new drivers to the fleet.",
                        parameters=genai.protos.Schema(
                            type=genai.protos.Type.OBJECT,
                            properties={
                                "name": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Full name of the driver"
                                ),
                                "phone": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Driver phone number"
                                ),
                                "email": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Driver email address (optional)"
                                ),
                                "vehicle_type": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Type of vehicle: van, truck, car, motorcycle (default: van)"
                                ),
                                "vehicle_plate": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Vehicle license plate number"
                                ),
                                "capacity_kg": genai.protos.Schema(
                                    type=genai.protos.Type.NUMBER,
                                    description="Vehicle cargo capacity in kilograms (default: 1000.0)"
                                ),
                                "capacity_m3": genai.protos.Schema(
                                    type=genai.protos.Type.NUMBER,
                                    description="Vehicle cargo volume in cubic meters (default: 12.0)"
                                ),
                                "skills": genai.protos.Schema(
                                    type=genai.protos.Type.ARRAY,
                                    description="List of driver skills/certifications: refrigerated, medical_certified, fragile_handler, overnight, express_delivery",
                                    items=genai.protos.Schema(type=genai.protos.Type.STRING)
                                ),
                                "status": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Driver status: active, busy, offline, unavailable (default: active)"
                                )
                            },
                            required=["name"]
                        )
                    ),
                    genai.protos.FunctionDeclaration(
                        name="count_orders",
                        description="Count total orders in the database with optional filters. Use this when user asks 'how many orders', 'fetch orders', or wants to know order statistics.",
                        parameters=genai.protos.Schema(
                            type=genai.protos.Type.OBJECT,
                            properties={
                                "status": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Filter by order status: pending, assigned, in_transit, delivered, failed, cancelled (optional)"
                                ),
                                "priority": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Filter by priority level: standard, express, urgent (optional)"
                                ),
                                "payment_status": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Filter by payment status: pending, paid, cod (optional)"
                                ),
                                "assigned_driver_id": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Filter by assigned driver ID (optional)"
                                ),
                                "is_fragile": genai.protos.Schema(
                                    type=genai.protos.Type.BOOLEAN,
                                    description="Filter fragile packages only (optional)"
                                ),
                                "requires_signature": genai.protos.Schema(
                                    type=genai.protos.Type.BOOLEAN,
                                    description="Filter orders requiring signature (optional)"
                                ),
                                "requires_cold_storage": genai.protos.Schema(
                                    type=genai.protos.Type.BOOLEAN,
                                    description="Filter orders requiring cold storage (optional)"
                                )
                            },
                            required=[]
                        )
                    ),
                    genai.protos.FunctionDeclaration(
                        name="fetch_orders",
                        description="Fetch orders from the database with optional filters, pagination, and sorting. Use after counting to show specific number of orders.",
                        parameters=genai.protos.Schema(
                            type=genai.protos.Type.OBJECT,
                            properties={
                                "limit": genai.protos.Schema(
                                    type=genai.protos.Type.INTEGER,
                                    description="Number of orders to fetch (default: 10, max: 100)"
                                ),
                                "offset": genai.protos.Schema(
                                    type=genai.protos.Type.INTEGER,
                                    description="Number of orders to skip for pagination (default: 0)"
                                ),
                                "status": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Filter by order status: pending, assigned, in_transit, delivered, failed, cancelled (optional)"
                                ),
                                "priority": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Filter by priority level: standard, express, urgent (optional)"
                                ),
                                "payment_status": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Filter by payment status: pending, paid, cod (optional)"
                                ),
                                "assigned_driver_id": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Filter by assigned driver ID (optional)"
                                ),
                                "is_fragile": genai.protos.Schema(
                                    type=genai.protos.Type.BOOLEAN,
                                    description="Filter fragile packages only (optional)"
                                ),
                                "requires_signature": genai.protos.Schema(
                                    type=genai.protos.Type.BOOLEAN,
                                    description="Filter orders requiring signature (optional)"
                                ),
                                "requires_cold_storage": genai.protos.Schema(
                                    type=genai.protos.Type.BOOLEAN,
                                    description="Filter orders requiring cold storage (optional)"
                                ),
                                "sort_by": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Field to sort by: created_at, priority, time_window_start (default: created_at)"
                                ),
                                "sort_order": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Sort order: ASC, DESC (default: DESC for newest first)"
                                )
                            },
                            required=[]
                        )
                    ),
                    genai.protos.FunctionDeclaration(
                        name="get_order_details",
                        description="Get complete details of a specific order by order ID. Use when user asks 'tell me about order X' or wants detailed information about a specific order.",
                        parameters=genai.protos.Schema(
                            type=genai.protos.Type.OBJECT,
                            properties={
                                "order_id": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="The order ID to fetch details for (e.g., 'ORD-20251114163800')"
                                )
                            },
                            required=["order_id"]
                        )
                    ),
                    genai.protos.FunctionDeclaration(
                        name="search_orders",
                        description="Search for orders by customer name, email, phone, or order ID pattern. Use when user provides partial information to find orders.",
                        parameters=genai.protos.Schema(
                            type=genai.protos.Type.OBJECT,
                            properties={
                                "search_term": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Search term to match against customer_name, customer_email, customer_phone, or order_id"
                                )
                            },
                            required=["search_term"]
                        )
                    ),
                    genai.protos.FunctionDeclaration(
                        name="get_incomplete_orders",
                        description="Get all orders that are not yet completed (excludes delivered and cancelled orders). Shortcut for finding orders in progress (pending, assigned, in_transit).",
                        parameters=genai.protos.Schema(
                            type=genai.protos.Type.OBJECT,
                            properties={
                                "limit": genai.protos.Schema(
                                    type=genai.protos.Type.INTEGER,
                                    description="Number of orders to fetch (default: 20)"
                                )
                            },
                            required=[]
                        )
                    ),
                    genai.protos.FunctionDeclaration(
                        name="count_drivers",
                        description="Count total drivers in the database with optional filters. Use this when user asks 'how many drivers', 'show drivers', or wants driver statistics.",
                        parameters=genai.protos.Schema(
                            type=genai.protos.Type.OBJECT,
                            properties={
                                "status": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Filter by driver status: active, busy, offline, unavailable (optional)"
                                ),
                                "vehicle_type": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Filter by vehicle type: van, truck, car, motorcycle, etc. (optional)"
                                )
                            },
                            required=[]
                        )
                    ),
                    genai.protos.FunctionDeclaration(
                        name="fetch_drivers",
                        description="Fetch drivers from the database with optional filters, pagination, and sorting. Use after counting to show specific number of drivers.",
                        parameters=genai.protos.Schema(
                            type=genai.protos.Type.OBJECT,
                            properties={
                                "limit": genai.protos.Schema(
                                    type=genai.protos.Type.INTEGER,
                                    description="Number of drivers to fetch (default: 10, max: 100)"
                                ),
                                "offset": genai.protos.Schema(
                                    type=genai.protos.Type.INTEGER,
                                    description="Number of drivers to skip for pagination (default: 0)"
                                ),
                                "status": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Filter by driver status: active, busy, offline, unavailable (optional)"
                                ),
                                "vehicle_type": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Filter by vehicle type: van, truck, car, motorcycle, etc. (optional)"
                                ),
                                "sort_by": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Field to sort by: name, status, created_at, last_location_update (default: name)"
                                ),
                                "sort_order": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Sort order: ASC, DESC (default: ASC for alphabetical)"
                                )
                            },
                            required=[]
                        )
                    ),
                    genai.protos.FunctionDeclaration(
                        name="get_driver_details",
                        description="Get complete details of a specific driver by driver ID, including current location (latitude, longitude, and human-readable address), contact info, vehicle details, status, and skills. Use when user asks about a driver's location, coordinates, position, or any other driver information.",
                        parameters=genai.protos.Schema(
                            type=genai.protos.Type.OBJECT,
                            properties={
                                "driver_id": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="The driver ID to fetch details for (e.g., 'DRV-20251114163800')"
                                )
                            },
                            required=["driver_id"]
                        )
                    ),
                    genai.protos.FunctionDeclaration(
                        name="search_drivers",
                        description="Search for drivers by name, email, phone, vehicle plate, or driver ID pattern. Use when user provides partial information to find drivers.",
                        parameters=genai.protos.Schema(
                            type=genai.protos.Type.OBJECT,
                            properties={
                                "search_term": genai.protos.Schema(
                                    type=genai.protos.Type.STRING,
                                    description="Search term to match against name, email, phone, vehicle_plate, or driver_id"
                                )
                            },
                            required=["search_term"]
                        )
                    ),
                    genai.protos.FunctionDeclaration(
                        name="get_available_drivers",
                        description="Get all drivers that are available for assignment (active or offline status, excludes busy and unavailable). Shortcut for finding drivers ready for dispatch.",
                        parameters=genai.protos.Schema(
                            type=genai.protos.Type.OBJECT,
                            properties={
                                "limit": genai.protos.Schema(
                                    type=genai.protos.Type.INTEGER,
                                    description="Number of drivers to fetch (default: 20)"
                                )
                            },
                            required=[]
                        )
                    ),
                    genai.protos.FunctionDeclaration(
                        name="update_order",
                        description="Update an existing order's details. You can update any combination of fields. Only provide the fields you want to change.",
                        parameters=genai.protos.Schema(
                            type=genai.protos.Type.OBJECT,
                            properties={
                                "order_id": genai.protos.Schema(type=genai.protos.Type.STRING, description="Order ID to update"),
                                "customer_name": genai.protos.Schema(type=genai.protos.Type.STRING, description="Updated customer name"),
                                "customer_phone": genai.protos.Schema(type=genai.protos.Type.STRING, description="Updated customer phone"),
                                "customer_email": genai.protos.Schema(type=genai.protos.Type.STRING, description="Updated customer email"),
                                "delivery_address": genai.protos.Schema(type=genai.protos.Type.STRING, description="Updated delivery address"),
                                "delivery_lat": genai.protos.Schema(type=genai.protos.Type.NUMBER, description="Updated delivery latitude"),
                                "delivery_lng": genai.protos.Schema(type=genai.protos.Type.NUMBER, description="Updated delivery longitude"),
                                "status": genai.protos.Schema(type=genai.protos.Type.STRING, description="Updated order status (pending/assigned/in_transit/delivered/failed/cancelled)"),
                                "priority": genai.protos.Schema(type=genai.protos.Type.STRING, description="Updated priority (standard/express/urgent)"),
                                "special_instructions": genai.protos.Schema(type=genai.protos.Type.STRING, description="Updated special instructions"),
                                "time_window_end": genai.protos.Schema(type=genai.protos.Type.STRING, description="Updated delivery deadline"),
                                "payment_status": genai.protos.Schema(type=genai.protos.Type.STRING, description="Updated payment status (pending/paid/cod)"),
                                "weight_kg": genai.protos.Schema(type=genai.protos.Type.NUMBER, description="Updated weight in kg"),
                                "order_value": genai.protos.Schema(type=genai.protos.Type.NUMBER, description="Updated order value")
                            },
                            required=["order_id"]
                        )
                    ),
                    genai.protos.FunctionDeclaration(
                        name="delete_order",
                        description="Permanently delete an order from the database. This action cannot be undone. Use with caution.",
                        parameters=genai.protos.Schema(
                            type=genai.protos.Type.OBJECT,
                            properties={
                                "order_id": genai.protos.Schema(type=genai.protos.Type.STRING, description="Order ID to delete"),
                                "confirm": genai.protos.Schema(type=genai.protos.Type.BOOLEAN, description="Must be true to confirm deletion")
                            },
                            required=["order_id", "confirm"]
                        )
                    ),
                    genai.protos.FunctionDeclaration(
                        name="update_driver",
                        description="Update an existing driver's details. You can update any combination of fields. Only provide the fields you want to change.",
                        parameters=genai.protos.Schema(
                            type=genai.protos.Type.OBJECT,
                            properties={
                                "driver_id": genai.protos.Schema(type=genai.protos.Type.STRING, description="Driver ID to update"),
                                "name": genai.protos.Schema(type=genai.protos.Type.STRING, description="Updated driver name"),
                                "phone": genai.protos.Schema(type=genai.protos.Type.STRING, description="Updated phone"),
                                "email": genai.protos.Schema(type=genai.protos.Type.STRING, description="Updated email"),
                                "status": genai.protos.Schema(type=genai.protos.Type.STRING, description="Updated status (active/busy/offline/unavailable)"),
                                "vehicle_type": genai.protos.Schema(type=genai.protos.Type.STRING, description="Updated vehicle type"),
                                "vehicle_plate": genai.protos.Schema(type=genai.protos.Type.STRING, description="Updated vehicle plate"),
                                "capacity_kg": genai.protos.Schema(type=genai.protos.Type.NUMBER, description="Updated capacity in kg"),
                                "capacity_m3": genai.protos.Schema(type=genai.protos.Type.NUMBER, description="Updated capacity in m³"),
                                "current_lat": genai.protos.Schema(type=genai.protos.Type.NUMBER, description="Updated current latitude"),
                                "current_lng": genai.protos.Schema(type=genai.protos.Type.NUMBER, description="Updated current longitude")
                            },
                            required=["driver_id"]
                        )
                    ),
                    genai.protos.FunctionDeclaration(
                        name="delete_driver",
                        description="Permanently delete a driver from the database. This action cannot be undone. Use with caution.",
                        parameters=genai.protos.Schema(
                            type=genai.protos.Type.OBJECT,
                            properties={
                                "driver_id": genai.protos.Schema(type=genai.protos.Type.STRING, description="Driver ID to delete"),
                                "confirm": genai.protos.Schema(type=genai.protos.Type.BOOLEAN, description="Must be true to confirm deletion")
                            },
                            required=["driver_id", "confirm"]
                        )
                    )
                ]
            )
        ]

    def _ensure_initialized(self):
        """Lazy initialization - only create model when first needed"""
        if self._initialized or not self.api_available:
            return

        try:
            genai.configure(api_key=self.api_key)
            self.model = genai.GenerativeModel(
                model_name=self.model_name,
                tools=self._get_gemini_tools(),
                system_instruction=self._get_system_prompt()
            )
            self._initialized = True
            logger.info(f"GeminiProvider: Model initialized ({self.model_name})")
        except Exception as e:
            logger.error(f"GeminiProvider: Failed to initialize: {e}")
            self.api_available = False
            self.model = None

    def is_available(self) -> bool:
        return self.api_available

    def get_status(self) -> str:
        if self.api_available:
            return f"✅ Connected - Model: {self.model_name}"
        return "⚠️ Not configured (add GOOGLE_API_KEY)"

    def get_provider_name(self) -> str:
        return "Gemini (Google)"

    def get_model_name(self) -> str:
        return self.model_name if self.api_available else "gemini-2.0-flash"

    def process_message(
        self,
        user_message: str,
        conversation
    ) -> Tuple[str, List[Dict]]:
        """Process user message with Gemini"""
        if not self.api_available:
            return self._handle_no_api(), []

        # Lazy initialization on first use
        self._ensure_initialized()

        if not self._initialized:
            return "⚠️ Failed to initialize Gemini model. Please check your API key and try again.", []

        try:
            # Build conversation history for Gemini
            chat = self.model.start_chat(history=self._convert_history(conversation))

            # Send message and get response
            response = chat.send_message(
                user_message,
                safety_settings={
                    HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
                    HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
                    HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
                    HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
                }
            )

            # Add user message to conversation
            conversation.add_message("user", user_message)

            # Process response and handle function calls
            return self._process_response(response, conversation, chat)

        except Exception as e:
            error_msg = f"⚠️ Gemini API error: {str(e)}"
            logger.error(f"Gemini provider error: {e}")
            return error_msg, []

    def _convert_history(self, conversation) -> list:
        """Convert conversation history to Gemini format"""
        history = []
        # Get all messages from conversation (history is built before adding current message)
        for msg in conversation.get_history():
            role = "user" if msg["role"] == "user" else "model"
            history.append({
                "role": role,
                "parts": [{"text": str(msg["content"])}]
            })
        return history

    def _process_response(
        self,
        response,
        conversation,
        chat
    ) -> Tuple[str, List[Dict]]:
        """Process Gemini's response and handle function calls"""
        tool_calls_made = []

        # Check if Gemini wants to call functions
        try:
            # Check ALL parts for function calls (not just first)
            has_function_call = False
            parts = response.candidates[0].content.parts
            logger.info(f"Processing response with {len(parts)} part(s)")

            for part in parts:
                if hasattr(part, 'function_call'):
                    fc = part.function_call
                    # More robust check
                    if fc is not None:
                        try:
                            if hasattr(fc, 'name') and fc.name:
                                has_function_call = True
                                logger.info(f"Detected function call: {fc.name}")
                                break
                        except Exception as e:
                            logger.warning(f"Error checking function call: {e}")

            if has_function_call:
                # Handle function calls (potentially multiple in sequence)
                current_response = response
                max_iterations = 3  # Prevent rate limit errors (Gemini free tier: 15 req/min)

                for iteration in range(max_iterations):
                    # Check if current response has a function call
                    try:
                        parts = current_response.candidates[0].content.parts
                        logger.info(f"Iteration {iteration + 1}: Response has {len(parts)} part(s)")
                    except (IndexError, AttributeError) as e:
                        logger.error(f"Cannot access response parts: {e}")
                        break

                    # Check ALL parts for function calls (some responses have text + function_call)
                    has_fc = False
                    fc_part = None

                    for idx, part in enumerate(parts):
                        if hasattr(part, 'function_call'):
                            fc = part.function_call
                            if fc and hasattr(fc, 'name') and fc.name:
                                has_fc = True
                                fc_part = part
                                logger.info(f"Iteration {iteration + 1}: Found function_call in part {idx}: {fc.name}")
                                break

                        # Also check if there's text (indicates Gemini wants to respond instead of continuing)
                        if hasattr(part, 'text') and part.text:
                            logger.warning(f"Iteration {iteration + 1}: Part {idx} has text: {part.text[:100]}")

                    if not has_fc:
                        # No more function calls, break and extract text
                        logger.info(f"No more function calls after iteration {iteration + 1}")
                        break

                    # Use the part with function_call
                    first_part = fc_part

                    # Extract function call details
                    function_call = first_part.function_call
                    function_name = function_call.name
                    function_args = dict(function_call.args) if function_call.args else {}

                    logger.info(f"Gemini executing function: {function_name} (iteration {iteration + 1})")

                    # Execute the tool
                    tool_result = execute_tool(function_name, function_args)

                    # Track for transparency
                    tool_calls_made.append({
                        "tool": function_name,
                        "input": function_args,
                        "result": tool_result
                    })

                    conversation.add_tool_result(function_name, function_args, tool_result)

                    # Send function result back to Gemini
                    try:
                        current_response = chat.send_message(
                            genai.protos.Content(
                                parts=[genai.protos.Part(
                                    function_response=genai.protos.FunctionResponse(
                                        name=function_name,
                                        response={"result": tool_result}
                                    )
                                )]
                            )
                        )
                    except Exception as e:
                        logger.error(f"Error sending function response: {e}")
                        break

                # Now extract text from the final response
                # NEVER use .text property directly - always check parts
                final_text = ""
                try:
                    parts = current_response.candidates[0].content.parts
                    logger.info(f"Extracting text from {len(parts)} parts")

                    for idx, part in enumerate(parts):
                        # Check if this part has a function call
                        if hasattr(part, 'function_call') and part.function_call:
                            fc = part.function_call
                            if hasattr(fc, 'name') and fc.name:
                                logger.warning(f"Part {idx} still has function call: {fc.name}. Skipping.")
                                continue

                        # Extract text from this part
                        if hasattr(part, 'text') and part.text:
                            logger.info(f"Part {idx} has text: {part.text[:50]}...")
                            final_text += part.text

                except (AttributeError, IndexError) as e:
                    logger.error(f"Error extracting text from parts: {e}")

                # Generate fallback message if still no text
                if not final_text:
                    logger.warning("No text extracted from response, generating fallback")
                    if tool_calls_made:
                        # Create a summary of what was done
                        tool_names = [t["tool"] for t in tool_calls_made]
                        if "create_order" in tool_names:
                            # Check if order was created successfully
                            create_result = next((t["result"] for t in tool_calls_made if t["tool"] == "create_order"), {})
                            if create_result.get("success"):
                                order_id = create_result.get("order_id", "")
                                final_text = f"✅ Order {order_id} created successfully!"
                            else:
                                final_text = "⚠️ There was an issue creating the order."
                        else:
                            final_text = f"✅ Executed {len(tool_calls_made)} tool(s) successfully!"
                    else:
                        final_text = "✅ Task completed!"

                logger.info(f"Returning response: {final_text[:100]}")
                conversation.add_message("assistant", final_text)
                return final_text, tool_calls_made

            else:
                # No function call detected, extract text from parts
                text_response = ""
                try:
                    parts = response.candidates[0].content.parts
                    logger.info(f"Extracting text from {len(parts)} parts (no function call)")

                    for idx, part in enumerate(parts):
                        # Double-check no function call in this part
                        if hasattr(part, 'function_call') and part.function_call:
                            fc = part.function_call
                            if hasattr(fc, 'name') and fc.name:
                                logger.error(f"Part {idx} has function call {fc.name} but was not detected earlier!")
                                # We missed a function call - handle it now
                                logger.info("Re-processing response with function call handling")
                                return self._process_response(response, conversation, chat)

                        # Extract text
                        if hasattr(part, 'text') and part.text:
                            logger.info(f"Part {idx} has text: {part.text[:50]}...")
                            text_response += part.text

                except (ValueError, AttributeError, IndexError) as e:
                    logger.error(f"Error extracting text from response: {e}")

                # Fallback if no text extracted
                if not text_response:
                    logger.warning("No text in response, using fallback")
                    text_response = "I'm ready to help! What would you like me to do?"

                conversation.add_message("assistant", text_response)
                return text_response, tool_calls_made

        except Exception as e:
            logger.error(f"Error processing Gemini response: {e}")
            error_msg = f"⚠️ Error processing response: {str(e)}"
            conversation.add_message("assistant", error_msg)
            return error_msg, tool_calls_made

    def _handle_no_api(self) -> str:
        """Return error message when API is not available"""
        return """⚠️ **Gemini API requires Google API key**

To use Gemini:

1. Get an API key from: https://aistudio.google.com/app/apikey
   - Free tier: 15 requests/min, 1500/day
   - Or use hackathon credits

2. Add to your `.env` file:
   ```
   GOOGLE_API_KEY=your-gemini-key-here
   ```

3. Restart the application

**Alternative:** Switch to Claude by setting `AI_PROVIDER=anthropic` in .env
"""

    def get_welcome_message(self) -> str:
        if not self.api_available:
            return self._handle_no_api()

        # Don't initialize model yet - wait for first actual message
        # This allows the welcome message to load instantly

        return """👋 Hello! I'm your AI dispatch assistant powered by **Google Gemini 2.0 Flash**.

I can help you manage your delivery fleet!

---

📋 **What I Can Do:**

**1. Create Delivery Orders:**
• Customer Name
• Delivery Address
• Contact (Phone OR Email)
• Optional: Deadline, Priority, Special Instructions

**2. Query & Search Orders:**
• Fetch orders with filters (status, priority, payment, etc.)
• Get complete details of specific orders
• Search by customer name, phone, email, or order ID
• Find incomplete/pending orders

**3. Create New Drivers:**
• Driver Name (required)
• Optional: Phone, Email, Vehicle Type, License Plate, Skills

**4. Query & Search Drivers:**
• Fetch drivers with filters (status, vehicle type)
• Get complete details of specific drivers
• Search by name, phone, email, plate, or driver ID
• Find available/active drivers

---

**Examples - Just Type Naturally:**

📦 **Create Orders:**
💬 "Create order for John Doe, 123 Main St San Francisco CA, phone 555-1234, deliver by 5 PM"
💬 "New urgent delivery to Sarah at 456 Oak Ave NYC, email [email protected]"

🔍 **Query Orders:**
💬 "Fetch the orders" or "Show me orders"
💬 "Which orders are incomplete?"
💬 "Tell me about order ORD-20251114120000"
💬 "Show me 10 urgent orders"
💬 "Search for orders from John"

🚚 **Create Drivers:**
💬 "Add new driver Tom Wilson, phone 555-0101, drives a van, plate ABC-123"
💬 "Create driver Sarah Martinez with refrigerated truck, phone 555-0202"
💬 "New driver: Mike Chen, email [email protected], motorcycle delivery"

👥 **Query Drivers:**
💬 "Show me drivers" or "Fetch the drivers"
💬 "Which drivers are available?"
💬 "Tell me about driver DRV-20251114163800"
💬 "Show 5 active drivers with vans"
💬 "Search for Tom"

---

🚀 **I'll automatically:**
• Geocode addresses for orders
• Generate unique IDs
• Save everything to the database
• Filter and search across all order fields

What would you like to do?"""