src/foodwise/agent/tools.py

"""
Inventory management tools for the FoodWise agent.

These tools provide the core functionality for interacting with the Notion-based
food inventory database through the LangGraph agent.
"""
from typing import List, Dict, Any, Optional, Union
from datetime import datetime, date, timedelta
from pydantic import BaseModel, Field
from langchain_core.tools import tool

from ..notion.notion_client import FoodWiseNotionClient


class InventoryItem(BaseModel):
    """Model for inventory items following the Notion database schema."""
    name: str = Field(description="Name of the food item")
    category: str = Field(description="Food category (Dairy, Raw Meat, Vegetables, etc.)")
    storage_type: str = Field(description="Storage type (Pantry, Refrigerator, Freezer)")
    quantity: float = Field(description="Quantity of the item")
    unit: str = Field(description="Unit of measurement (lbs, oz, pieces, etc.)")
    best_by_date: Optional[str] = Field(default=None, description="Best by/expiration date (YYYY-MM-DD)")
    purchase_date: Optional[str] = Field(default=None, description="Purchase date (YYYY-MM-DD)")
    opened_date: Optional[str] = Field(default=None, description="Date opened (YYYY-MM-DD)")
    location_shelf: Optional[str] = Field(default=None, description="Specific storage location/shelf")
    fridge_zone: Optional[str] = Field(default=None, description="Fridge zone for refrigerated items")
    tags: Optional[List[str]] = Field(default=None, description="Tags (Low Stock, Organic, etc.)")
    temperature_sensitive: bool = Field(default=False, description="Requires temperature control")
    cooked_raw_status: str = Field(default="Raw", description="Cooking status (Raw, Cooked, Partially Cooked)")
    prep_notes: Optional[str] = Field(default=None, description="Preparation notes")


@tool
def get_inventory(
    filter_type: Optional[str] = None,
    category: Optional[str] = None,
    storage_type: Optional[str] = None,
    expiring_days: Optional[int] = None
) -> List[Dict[str, Any]]:
    """
    Retrieve inventory items with optional filtering.
    
    Args:
        filter_type: Filter by type ("expiring_soon", "low_stock", "all")
        category: Filter by specific category
        storage_type: Filter by storage type (Pantry, Refrigerator, Freezer)
        expiring_days: Show items expiring within this many days
        
    Returns:
        List of inventory items matching the criteria
    """
    try:
        notion_client = FoodWiseNotionClient()
        
        # Handle different filter types
        if filter_type == "expiring_soon" or expiring_days:
            days = expiring_days or 7
            return notion_client.get_expiring_items(days)
        
        # Build filter conditions
        filter_conditions = None
        if category or storage_type:
            conditions = []
            
            if category:
                conditions.append({
                    "property": "Category",
                    "select": {"equals": category}
                })
            
            if storage_type:
                conditions.append({
                    "property": "Storage Type", 
                    "select": {"equals": storage_type}
                })
            
            if len(conditions) == 1:
                filter_conditions = conditions[0]
            else:
                filter_conditions = {"and": conditions}
        
        return notion_client.query_inventory(filter_conditions=filter_conditions)
        
    except Exception as e:
        # Fallback to sample data if Notion fails
        return [{"error": f"Failed to connect to Notion: {e}"}]


@tool 
def add_item(
    name: str,
    category: str,
    storage_type: str,
    quantity: float,
    unit: str,
    best_by_date: Optional[str] = None,
    purchase_date: Optional[str] = None,
    location_shelf: Optional[str] = None,
    fridge_zone: Optional[str] = None,
    tags: Optional[List[str]] = None,
    temperature_sensitive: bool = False,
    cooked_raw_status: str = "Raw",
    prep_notes: Optional[str] = None
) -> Dict[str, Any]:
    """
    Add a new food item to the inventory.
    
    Args:
        name: Name of the food item
        category: Food category (must match schema options)
        storage_type: Storage type (Pantry, Refrigerator, Freezer)
        quantity: Quantity of the item
        unit: Unit of measurement
        best_by_date: Best by date (YYYY-MM-DD format)
        purchase_date: Purchase date (YYYY-MM-DD format)
        location_shelf: Specific storage location
        fridge_zone: Fridge zone for refrigerated items
        tags: List of tags
        temperature_sensitive: Whether item requires temperature control
        cooked_raw_status: Cooking status
        prep_notes: Preparation notes
        
    Returns:
        Success status and item details
    """
    # Validate required fields
    valid_categories = [
        "Dairy", "Raw Meat", "Cooked Meat", "Vegetables", "Fruits", 
        "Canned Goods", "Grains", "Pasta", "Spices", "Oils",
        "Baking Supplies", "Snacks", "Condiments", "Leftovers", 
        "Beverages", "Prepared Meals", "Ice Cream"
    ]
    
    if category not in valid_categories:
        return {"success": False, "error": f"Invalid category. Must be one of: {valid_categories}"}
    
    valid_storage_types = ["Pantry", "Refrigerator", "Freezer"]
    if storage_type not in valid_storage_types:
        return {"success": False, "error": f"Invalid storage type. Must be one of: {valid_storage_types}"}
    
    try:
        notion_client = FoodWiseNotionClient()
        
        # Prepare item data
        item_data = {
            "name": name,
            "category": category,
            "storage_type": storage_type,
            "quantity": quantity,
            "unit": unit,
            "best_by_date": best_by_date,
            "purchase_date": purchase_date,
            "location_shelf": location_shelf,
            "fridge_zone": fridge_zone,
            "tags": tags,
            "temperature_sensitive": temperature_sensitive,
            "cooked_raw_status": cooked_raw_status,
            "prep_notes": prep_notes
        }
        
        # Add to Notion database
        created_item = notion_client.add_inventory_item(item_data)
        
        return {
            "success": True,
            "message": f"Successfully added {quantity} {unit} of {name} to {storage_type}",
            "item": created_item
        }
        
    except Exception as e:
        return {
            "success": False,
            "error": f"Failed to add item to Notion: {e}"
        }


@tool
def query_item(
    search_term: str,
    search_type: str = "name"
) -> List[Dict[str, Any]]:
    """
    Search for specific items in the inventory.
    
    Args:
        search_term: Term to search for
        search_type: Type of search ("name", "category", "tag")
        
    Returns:
        List of items matching the search criteria
    """
    try:
        notion_client = FoodWiseNotionClient()
        
        if search_type == "name":
            return notion_client.search_items(search_term)
        elif search_type == "category":
            return notion_client.query_inventory(filter_conditions={
                "property": "Category",
                "select": {"equals": search_term}
            })
        else:
            # Fallback to name search
            return notion_client.search_items(search_term)
            
    except Exception as e:
        return [{"error": f"Failed to search Notion: {e}"}]


class UpdateItemInput(BaseModel):
    page_id: str = Field(description="Notion page ID of the item to update")
    updated_data: Dict[str, Any] = Field(description="Dictionary of fields to update")

@tool(args_schema=UpdateItemInput)
def update_item(page_id: str, updated_data: Dict[str, Any]) -> Dict[str, Any]:
    """Update an existing inventory item."""
    try:
        notion_client = FoodWiseNotionClient()
        updated_item = notion_client.update_inventory_item(page_id, updated_data)
        return {
            "success": True,
            "message": f"Item {page_id} updated successfully",
            "item": updated_item
        }
    except Exception as e:
        return {"success": False, "error": f"Failed to update item: {str(e)}"}

class RemoveItemInput(BaseModel):
    page_id: str = Field(description="Notion page ID of the item to remove")

@tool(args_schema=RemoveItemInput)
def remove_item(page_id: str) -> Dict[str, Any]:
    """Remove an inventory item."""
    try:
        notion_client = FoodWiseNotionClient()
        result = notion_client.remove_inventory_item(page_id)
        return result
    except Exception as e:
        return {"success": False, "error": f"Failed to remove item: {str(e)}"}


def get_inventory_tools() -> List:
    """
    Get all available inventory management tools.
    
    Returns:
        List of LangChain tools for the agent
    """
    return [
        get_inventory,
        add_item,
        query_item,
        update_item,
        remove_item
    ]