src/wanderlust_ai/utils/notebook_utils.py

"""
Notebook utility functions for development and testing.

This module provides utilities for setting up notebook environments,
displaying results, and common development tasks.
"""

import sys
import os
import time
from pathlib import Path
from typing import Dict, Any, Optional
from datetime import datetime

from IPython.display import display, JSON

from .config import get_settings


def setup_notebook_environment(project_root: Optional[Path] = None) -> Dict[str, Any]:
    """
    Setup notebook environment with proper configuration.
    
    This function replaces the need for %run commands by providing
    a clean, explicit way to setup the notebook environment.
    
    Args:
        project_root: Optional path to project root. Auto-detects if None.
        
    Returns:
        Dictionary with setup information including settings and available modules
    """
    # Auto-detect project root if not provided
    if project_root is None:
        current_dir = Path.cwd()
        if (current_dir / "src").exists():
            project_root = current_dir
        elif (current_dir.parent / "src").exists():
            project_root = current_dir.parent
        else:
            raise ValueError("Could not auto-detect project root. Please specify project_root.")
    
    # Add src directory to Python path
    src_dir = project_root / "src"
    if str(src_dir) not in sys.path:
        sys.path.insert(0, str(src_dir))
        print(f"✅ Added {src_dir} to Python path")
    else:
        print(f"✅ {src_dir} already in Python path")
    
    # Load settings
    try:
        settings = get_settings()
        print(f"✅ Configuration loaded successfully")
        print(f"   Debug mode: {settings.debug}")
        print(f"   Log level: {settings.log_level}")
        print(f"   Anthropic API configured: {'✅' if settings.anthropic_api_key else '❌'}")
        print(f"   Tavily API configured: {'✅' if settings.tavily_api_key else '❌'}")
    except Exception as e:
        print(f"❌ Configuration error: {e}")
        settings = None
    
    # Test module availability
    available_modules = {
        "anthropic": False,
        "pandas": False,
        "matplotlib": False,
        "seaborn": False,
        "httpx": False
    }
    
    # Test Anthropic
    try:
        from anthropic import AsyncAnthropic, Anthropic
        available_modules["anthropic"] = True
        print("✅ Anthropic imported successfully")
    except ImportError:
        print("⚠️  Anthropic not available - install with: pip install anthropic")
    
    # Test pandas
    try:
        import pandas as pd
        available_modules["pandas"] = True
        print("✅ pandas imported successfully")
    except ImportError:
        print("⚠️  pandas not available - install with: pip install pandas")
    
    # Test matplotlib
    try:
        import matplotlib.pyplot as plt
        available_modules["matplotlib"] = True
        print("✅ matplotlib imported successfully")
    except ImportError:
        print("⚠️  matplotlib not available - install with: pip install matplotlib")
    
    # Test seaborn
    try:
        import seaborn as sns
        available_modules["seaborn"] = True
        print("✅ seaborn imported successfully")
    except ImportError:
        print("⚠️  seaborn not available - install with: pip install seaborn")
    
    # Test httpx
    try:
        import httpx
        available_modules["httpx"] = True
        print("✅ httpx imported successfully")
    except ImportError:
        print("⚠️  httpx not available - install with: pip install httpx")
    
    print("✅ Core imports loaded successfully")
    
    return {
        "settings": settings,
        "project_root": project_root,
        "src_dir": src_dir,
        "available_modules": available_modules
    }


def display_json(data: Dict[str, Any], title: str = "JSON Data") -> None:
    """Display JSON data in a formatted way."""
    print(f"\n📋 {title}")
    print("=" * (len(title) + 4))
    display(JSON(data, expanded=True))


def display_api_response(response, title: str = "API Response") -> None:
    """Display API response in a formatted way."""
    print(f"\n🌐 {title}")
    print("=" * (len(title) + 4))
    
    if hasattr(response, 'success'):
        status = "✅ Success" if response.success else "❌ Failed"
        print(f"Status: {status}")
        print(f"Response time: {response.response_time_ms:.2f}ms")
        
        if response.success and response.data:
            display_json(response.data, "Response Data")
        elif response.error:
            print(f"Error: {response.error.message}")
            print(f"Error type: {response.error.error_type}")
    else:
        display_json(response, "Response Data")


def time_function(func, *args, **kwargs):
    """Time a function execution."""
    start_time = time.time()
    result = func(*args, **kwargs)
    end_time = time.time()
    
    print(f"⏱️  Function '{func.__name__}' took {end_time - start_time:.3f} seconds")
    return result


async def time_async_function(func, *args, **kwargs):
    """Time an async function execution."""
    start_time = time.time()
    result = await func(*args, **kwargs)
    end_time = time.time()
    
    print(f"⏱️  Async function '{func.__name__}' took {end_time - start_time:.3f} seconds")
    return result


def create_test_data() -> Dict[str, Any]:
    """Create sample test data for development."""
    return {
        "flight_request": {
            "departure_city": "New York",
            "arrival_city": "Los Angeles",
            "departure_date": "2024-01-15T08:30:00Z",
            "passengers": 2,
            "travel_class": "Economy"
        },
        "hotel_request": {
            "city": "Los Angeles",
            "check_in_date": "2024-01-15T15:00:00Z",
            "check_out_date": "2024-01-20T11:00:00Z",
            "guests": 2,
            "rooms": 1
        },
        "activity_request": {
            "city": "Los Angeles",
            "activity_date": "2024-01-16T10:00:00Z",
            "activity_type": "Sightseeing",
            "participants": 2
        }
    }


def show_settings_info(settings) -> None:
    """Display information about the settings object."""
    print("📋 Settings Information:")
    print("=" * 30)
    
    # Show available attributes
    attrs = [attr for attr in dir(settings) if not attr.startswith('_')]
    print(f"Available attributes: {len(attrs)}")
    
    # Show key configuration values
    key_attrs = ['debug', 'log_level', 'max_concurrent_requests', 'request_timeout']
    for attr in key_attrs:
        if hasattr(settings, attr):
            value = getattr(settings, attr)
            print(f"   {attr}: {value}")
    
    # Show API key status
    api_keys = ['anthropic_api_key', 'tavily_api_key', 'openai_api_key']
    print(f"\n🔑 API Key Status:")
    for key in api_keys:
        if hasattr(settings, key):
            value = getattr(settings, key)
            status = "✅ Configured" if value else "❌ Not configured"
            print(f"   {key}: {status}")


def test_notebook_setup() -> bool:
    """Test that the notebook setup is working correctly."""
    print("🧪 Testing notebook setup...")
    
    try:
        # Test configuration
        settings = get_settings()
        assert settings is not None, "Settings not loaded"
        print("✅ Configuration test passed")
        
        # Test imports
        from wanderlust_ai.api.travel_search_strategies import TravelSearchQueryBuilder
        query_builder = TravelSearchQueryBuilder()
        print("✅ Import test passed")
        
        # Test functionality
        queries = query_builder.build_flight_query("NYC", "LAX", "2024-01-15")
        assert len(queries['serpapi']) > 0, "No SerpAPI queries generated"
        assert len(queries['tavily']) > 0, "No Tavily queries generated"
        print("✅ Functionality test passed")
        
        print("🎉 All tests passed! Notebook setup is working correctly.")
        return True
        
    except Exception as e:
        print(f"❌ Test failed: {e}")
        return False