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:**
- geocode_address: Convert address to GPS coordinates
- create_order: Create customer delivery order (REQUIRES geocoded address)
- create_driver: Add new driver/delivery man to fleet

**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)

**Your goal:** Execute tasks in ONE smooth automated flow. No stopping, no intermediate messages!"""

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

    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 = 10  # Allow more iterations for complex tasks

                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()

        # Initialize on first use (welcome message)
        self._ensure_initialized()

        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. Add New Drivers:**
• Driver Name (required)
• Optional: Phone, Email, Vehicle Type, License Plate, Skills

---

**Examples - Just Type Naturally:**

📦 **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]"

🚚 **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"

---

🚀 **I'll automatically:**
• Geocode addresses for orders
• Generate unique IDs
• Save everything to the database

What would you like to do?"""