#!/usr/bin/env python3
"""
Script to create comprehensive test data for the Grocery Tracker application.
This includes shops, groceries, and shopping events with realistic data.
"""

import requests
import json
import random
import argparse
import sys
from datetime import datetime, timedelta
from typing import List, Dict, Any

BASE_URL = "http://localhost:8000"

# Test data definitions
SHOPS_DATA = [
    {"name": "Whole Foods Market", "city": "San Francisco", "address": "1765 California St"},
    {"name": "Safeway", "city": "San Francisco", "address": "2020 Market St"},
    {"name": "Trader Joe's", "city": "Berkeley", "address": "1885 University Ave"},
    {"name": "Berkeley Bowl", "city": "Berkeley", "address": "2020 Oregon St"},
    {"name": "Rainbow Grocery", "city": "San Francisco", "address": "1745 Folsom St"},
    {"name": "Mollie Stone's Market", "city": "Palo Alto", "address": "164 S California Ave"},
    {"name": "Costco Wholesale", "city": "San Mateo", "address": "2300 S Norfolk St"},
    {"name": "Target", "city": "Mountain View", "address": "1200 El Camino Real"},
    {"name": "Sprouts Farmers Market", "city": "Sunnyvale", "address": "1077 E El Camino Real"},
    {"name": "Lucky Supermarket", "city": "San Jose", "address": "1717 Tully Rd"},
]

GROCERIES_DATA = [
    # Fruits
    {"name": "Organic Bananas", "category": "Fruits", "organic": True, "weight": 1.0, "weight_unit": "lb"},
    {"name": "Gala Apples", "category": "Fruits", "organic": False, "weight": 2.0, "weight_unit": "lb"},
    {"name": "Organic Strawberries", "category": "Fruits", "organic": True, "weight": 1.0, "weight_unit": "lb"},
    {"name": "Avocados", "category": "Fruits", "organic": False, "weight": None, "weight_unit": "piece"},
    {"name": "Organic Blueberries", "category": "Fruits", "organic": True, "weight": 0.5, "weight_unit": "lb"},
    {"name": "Lemons", "category": "Fruits", "organic": False, "weight": None, "weight_unit": "piece"},
    {"name": "Organic Oranges", "category": "Fruits", "organic": True, "weight": 3.0, "weight_unit": "lb"},
    {"name": "Grapes", "category": "Fruits", "organic": False, "weight": 2.0, "weight_unit": "lb"},
    {"name": "Organic Pears", "category": "Fruits", "organic": True, "weight": 2.0, "weight_unit": "lb"},
    {"name": "Pineapple", "category": "Fruits", "organic": False, "weight": None, "weight_unit": "piece"},
    
    # Vegetables
    {"name": "Organic Spinach", "category": "Vegetables", "organic": True, "weight": 5.0, "weight_unit": "oz"},
    {"name": "Carrots", "category": "Vegetables", "organic": False, "weight": 2.0, "weight_unit": "lb"},
    {"name": "Organic Broccoli", "category": "Vegetables", "organic": True, "weight": None, "weight_unit": "piece"},
    {"name": "Red Bell Peppers", "category": "Vegetables", "organic": False, "weight": None, "weight_unit": "piece"},
    {"name": "Organic Kale", "category": "Vegetables", "organic": True, "weight": 1.0, "weight_unit": "bunch"},
    {"name": "Tomatoes", "category": "Vegetables", "organic": False, "weight": 2.0, "weight_unit": "lb"},
    {"name": "Organic Sweet Potatoes", "category": "Vegetables", "organic": True, "weight": 3.0, "weight_unit": "lb"},
    {"name": "Cucumbers", "category": "Vegetables", "organic": False, "weight": None, "weight_unit": "piece"},
    {"name": "Organic Lettuce", "category": "Vegetables", "organic": True, "weight": None, "weight_unit": "head"},
    {"name": "Onions", "category": "Vegetables", "organic": False, "weight": 3.0, "weight_unit": "lb"},
    
    # Dairy
    {"name": "Organic Whole Milk", "category": "Dairy", "organic": True, "weight": 1.0, "weight_unit": "gallon"},
    {"name": "Greek Yogurt", "category": "Dairy", "organic": False, "weight": 32.0, "weight_unit": "oz"},
    {"name": "Organic Eggs", "category": "Dairy", "organic": True, "weight": None, "weight_unit": "dozen"},
    {"name": "Cheddar Cheese", "category": "Dairy", "organic": False, "weight": 8.0, "weight_unit": "oz"},
    {"name": "Organic Butter", "category": "Dairy", "organic": True, "weight": 1.0, "weight_unit": "lb"},
    {"name": "Cream Cheese", "category": "Dairy", "organic": False, "weight": 8.0, "weight_unit": "oz"},
    {"name": "Organic Yogurt", "category": "Dairy", "organic": True, "weight": 6.0, "weight_unit": "oz"},
    
    # Meat & Seafood
    {"name": "Organic Chicken Breast", "category": "Meat & Seafood", "organic": True, "weight": 2.0, "weight_unit": "lb"},
    {"name": "Ground Beef", "category": "Meat & Seafood", "organic": False, "weight": 1.0, "weight_unit": "lb"},
    {"name": "Wild Salmon Fillet", "category": "Meat & Seafood", "organic": False, "weight": 1.5, "weight_unit": "lb"},
    {"name": "Organic Ground Turkey", "category": "Meat & Seafood", "organic": True, "weight": 1.0, "weight_unit": "lb"},
    {"name": "Shrimp", "category": "Meat & Seafood", "organic": False, "weight": 1.0, "weight_unit": "lb"},
    {"name": "Organic Chicken Thighs", "category": "Meat & Seafood", "organic": True, "weight": 2.0, "weight_unit": "lb"},
    
    # Pantry
    {"name": "Organic Brown Rice", "category": "Pantry", "organic": True, "weight": 2.0, "weight_unit": "lb"},
    {"name": "Whole Wheat Bread", "category": "Pantry", "organic": False, "weight": 24.0, "weight_unit": "oz"},
    {"name": "Organic Quinoa", "category": "Pantry", "organic": True, "weight": 1.0, "weight_unit": "lb"},
    {"name": "Olive Oil", "category": "Pantry", "organic": False, "weight": 500.0, "weight_unit": "ml"},
    {"name": "Organic Pasta", "category": "Pantry", "organic": True, "weight": 1.0, "weight_unit": "lb"},
    {"name": "Black Beans", "category": "Pantry", "organic": False, "weight": 15.0, "weight_unit": "oz"},
    {"name": "Organic Oats", "category": "Pantry", "organic": True, "weight": 18.0, "weight_unit": "oz"},
    {"name": "Peanut Butter", "category": "Pantry", "organic": False, "weight": 18.0, "weight_unit": "oz"},
    {"name": "Organic Honey", "category": "Pantry", "organic": True, "weight": 12.0, "weight_unit": "oz"},
    {"name": "Canned Tomatoes", "category": "Pantry", "organic": False, "weight": 14.5, "weight_unit": "oz"},
    
    # Beverages
    {"name": "Organic Orange Juice", "category": "Beverages", "organic": True, "weight": 64.0, "weight_unit": "oz"},
    {"name": "Sparkling Water", "category": "Beverages", "organic": False, "weight": 1.0, "weight_unit": "l"},
    {"name": "Organic Green Tea", "category": "Beverages", "organic": True, "weight": None, "weight_unit": "box"},
    {"name": "Coffee Beans", "category": "Beverages", "organic": False, "weight": 12.0, "weight_unit": "oz"},
    {"name": "Almond Milk", "category": "Beverages", "organic": False, "weight": 32.0, "weight_unit": "oz"},
    {"name": "Organic Apple Juice", "category": "Beverages", "organic": True, "weight": 64.0, "weight_unit": "oz"},
    
    # Frozen
    {"name": "Organic Frozen Berries", "category": "Frozen", "organic": True, "weight": 10.0, "weight_unit": "oz"},
    {"name": "Frozen Pizza", "category": "Frozen", "organic": False, "weight": 12.0, "weight_unit": "oz"},
    {"name": "Organic Frozen Vegetables", "category": "Frozen", "organic": True, "weight": 16.0, "weight_unit": "oz"},
    {"name": "Ice Cream", "category": "Frozen", "organic": False, "weight": 48.0, "weight_unit": "oz"},
    {"name": "Frozen Fish Fillets", "category": "Frozen", "organic": False, "weight": 1.0, "weight_unit": "lb"},
    
    # Snacks
    {"name": "Organic Granola Bars", "category": "Snacks", "organic": True, "weight": 8.0, "weight_unit": "oz"},
    {"name": "Potato Chips", "category": "Snacks", "organic": False, "weight": 5.0, "weight_unit": "oz"},
    {"name": "Organic Nuts", "category": "Snacks", "organic": True, "weight": 6.0, "weight_unit": "oz"},
    {"name": "Crackers", "category": "Snacks", "organic": False, "weight": 7.0, "weight_unit": "oz"},
    {"name": "Organic Popcorn", "category": "Snacks", "organic": True, "weight": 3.0, "weight_unit": "oz"},
]

# Price ranges for different categories (min, max)
PRICE_RANGES = {
    "Fruits": (1.99, 8.99),
    "Vegetables": (0.99, 6.99),
    "Dairy": (2.49, 12.99),
    "Meat & Seafood": (4.99, 24.99),
    "Pantry": (1.99, 15.99),
    "Beverages": (1.99, 8.99),
    "Frozen": (2.99, 9.99),
    "Snacks": (1.49, 7.99),
}

def parse_arguments():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description='Create test data for Grocery Tracker')
    parser.add_argument('--events', type=int, default=30, help='Number of shopping events to create (default: 30)')
    parser.add_argument('--days', type=int, default=90, help='Number of days back to generate events (default: 90)')
    parser.add_argument('--url', type=str, default=BASE_URL, help='API base URL (default: http://localhost:8000)')
    parser.add_argument('--shops-only', action='store_true', help='Create only shops')
    parser.add_argument('--groceries-only', action='store_true', help='Create only groceries')
    parser.add_argument('--events-only', action='store_true', help='Create only shopping events (requires existing shops and groceries)')
    parser.add_argument('--verbose', '-v', action='store_true', help='Verbose output')
    parser.add_argument('--dry-run', action='store_true', help='Show what would be created without actually creating it')
    return parser.parse_args()

def check_api_connection(base_url: str) -> bool:
    """Check if the API server is accessible."""
    try:
        response = requests.get(f"{base_url}/", timeout=5)
        return response.status_code == 200
    except requests.exceptions.RequestException:
        return False

def create_shops(base_url: str, verbose: bool = False, dry_run: bool = False) -> List[Dict[str, Any]]:
    """Create shops and return the created shop objects."""
    print("🏪 Creating shops...")
    created_shops = []
    
    if dry_run:
        print("   [DRY RUN] Would create the following shops:")
        for shop_data in SHOPS_DATA:
            print(f"   📋 {shop_data['name']} in {shop_data['city']}")
        return []
    
    for shop_data in SHOPS_DATA:
        try:
            if verbose:
                print(f"   🔄 Creating shop: {shop_data['name']}...")
            
            response = requests.post(f"{base_url}/shops/", json=shop_data, timeout=10)
            if response.status_code == 200:
                shop = response.json()
                created_shops.append(shop)
                print(f"   ✅ Created shop: {shop['name']} in {shop['city']}")
            else:
                print(f"   ❌ Failed to create shop {shop_data['name']}: {response.status_code}")
                if verbose:
                    print(f"      Response: {response.text}")
        except requests.exceptions.RequestException as e:
            print(f"   ❌ Network error creating shop {shop_data['name']}: {e}")
        except Exception as e:
            print(f"   ❌ Error creating shop {shop_data['name']}: {e}")
    
    print(f"   📊 Created {len(created_shops)} shops total\n")
    return created_shops

def create_groceries(base_url: str, verbose: bool = False, dry_run: bool = False) -> List[Dict[str, Any]]:
    """Create groceries and return the created grocery objects."""
    print("🥬 Creating groceries...")
    created_groceries = []
    
    if dry_run:
        print("   [DRY RUN] Would create the following groceries:")
        for grocery_data in GROCERIES_DATA:
            organic_label = "🌱" if grocery_data['organic'] else "🌾"
            print(f"   📋 {organic_label} {grocery_data['name']} ({grocery_data['category']})")
        return []
    
    for grocery_data in GROCERIES_DATA:
        try:
            if verbose:
                print(f"   🔄 Creating grocery: {grocery_data['name']}...")
            
            response = requests.post(f"{base_url}/groceries/", json=grocery_data, timeout=10)
            if response.status_code == 200:
                grocery = response.json()
                created_groceries.append(grocery)
                organic_label = "🌱" if grocery['organic'] else "🌾"
                print(f"   ✅ Created grocery: {organic_label} {grocery['name']} ({grocery['category']})")
            else:
                print(f"   ❌ Failed to create grocery {grocery_data['name']}: {response.status_code}")
                if verbose:
                    print(f"      Response: {response.text}")
        except requests.exceptions.RequestException as e:
            print(f"   ❌ Network error creating grocery {grocery_data['name']}: {e}")
        except Exception as e:
            print(f"   ❌ Error creating grocery {grocery_data['name']}: {e}")
    
    print(f"   📊 Created {len(created_groceries)} groceries total\n")
    return created_groceries

def generate_random_price(category: str, organic: bool = False) -> float:
    """Generate a random price for a grocery item based on category and organic status."""
    min_price, max_price = PRICE_RANGES.get(category, (1.99, 9.99))
    
    # Organic items are typically 20-50% more expensive
    if organic:
        min_price *= 1.2
        max_price *= 1.5
    
    # Generate random price and round to nearest cent
    price = random.uniform(min_price, max_price)
    return round(price, 2)

def get_existing_data(base_url: str) -> tuple[List[Dict], List[Dict]]:
    """Get existing shops and groceries from the API."""
    try:
        shops_response = requests.get(f"{base_url}/shops/", timeout=10)
        groceries_response = requests.get(f"{base_url}/groceries/", timeout=10)
        
        shops = shops_response.json() if shops_response.status_code == 200 else []
        groceries = groceries_response.json() if groceries_response.status_code == 200 else []
        
        return shops, groceries
    except requests.exceptions.RequestException as e:
        print(f"   ❌ Error fetching existing data: {e}")
        return [], []

def create_shopping_events(shops: List[Dict], groceries: List[Dict], base_url: str, 
                         num_events: int = 25, days_back: int = 90, 
                         verbose: bool = False, dry_run: bool = False) -> List[Dict[str, Any]]:
    """Create shopping events with random groceries and realistic data."""
    print(f"🛒 Creating {num_events} shopping events...")
    created_events = []
    
    if not shops:
        print("   ❌ No shops available. Cannot create shopping events.")
        return []
    
    if not groceries:
        print("   ❌ No groceries available. Cannot create shopping events.")
        return []
    
    # Generate events over the specified time period
    end_date = datetime.now()
    start_date = end_date - timedelta(days=days_back)
    
    if dry_run:
        print(f"   [DRY RUN] Would create {num_events} shopping events over {days_back} days")
        print(f"   📋 Date range: {start_date.strftime('%Y-%m-%d')} to {end_date.strftime('%Y-%m-%d')}")
        print(f"   📋 Available shops: {len(shops)}")
        print(f"   📋 Available groceries: {len(groceries)}")
        return []
    
    for i in range(num_events):
        try:
            if verbose:
                print(f"   🔄 Creating shopping event {i+1}/{num_events}...")
            
            # Random shop and date
            shop = random.choice(shops)
            event_date = start_date + timedelta(
                days=random.randint(0, days_back),
                hours=random.randint(8, 20),
                minutes=random.randint(0, 59)
            )
            
            # Random number of groceries (2-8 items per shopping trip)
            num_groceries = random.randint(2, 8)
            selected_groceries = random.sample(groceries, min(num_groceries, len(groceries)))
            
            # Create grocery items for this event
            event_groceries = []
            total_amount = 0.0
            
            for grocery in selected_groceries:
                # Random amount based on item type
                if grocery['weight_unit'] == 'piece':
                    amount = random.randint(1, 4)
                elif grocery['weight_unit'] == 'dozen':
                    amount = 1
                elif grocery['weight_unit'] in ['box', 'head', 'bunch']:
                    amount = random.randint(1, 2)
                elif grocery['weight_unit'] in ['gallon', 'l']:
                    amount = 1
                else:
                    amount = round(random.uniform(0.5, 3.0), 2)
                
                # Generate price based on category and organic status
                price = generate_random_price(grocery['category'], grocery['organic'])
                
                event_groceries.append({
                    "grocery_id": grocery['id'],
                    "amount": amount,
                    "price": price
                })
                
                total_amount += amount * price
            
            # Round total amount
            total_amount = round(total_amount, 2)
            
            # Random notes (30% chance of having notes)
            notes = None
            if random.random() < 0.3:
                note_options = [
                    "Weekly grocery shopping",
                    "Quick lunch ingredients",
                    "Dinner party prep",
                    "Meal prep for the week",
                    "Emergency grocery run",
                    "Organic produce haul",
                    "Bulk shopping trip",
                    "Special occasion shopping",
                    "Holiday meal preparation",
                    "Healthy eating restart",
                    "Stocking up on essentials",
                    "Trying new recipes",
                ]
                notes = random.choice(note_options)
            
            # Create the shopping event
            event_data = {
                "shop_id": shop['id'],
                "date": event_date.isoformat(),
                "total_amount": total_amount,
                "notes": notes,
                "groceries": event_groceries
            }
            
            response = requests.post(f"{base_url}/shopping-events/", json=event_data, timeout=15)
            if response.status_code == 200:
                event = response.json()
                created_events.append(event)
                print(f"   ✅ Created event #{event['id']}: {shop['name']} - ${total_amount:.2f} ({len(event_groceries)} items)")
            else:
                print(f"   ❌ Failed to create shopping event: {response.status_code}")
                if verbose:
                    print(f"      Response: {response.text}")
        except requests.exceptions.RequestException as e:
            print(f"   ❌ Network error creating shopping event {i+1}: {e}")
        except Exception as e:
            print(f"   ❌ Error creating shopping event {i+1}: {e}")
    
    print(f"   📊 Created {len(created_events)} shopping events total\n")
    return created_events

def print_summary(shops: List[Dict], groceries: List[Dict], events: List[Dict]):
    """Print a summary of the created test data."""
    print("📋 TEST DATA SUMMARY")
    print("=" * 50)
    
    print(f"🏪 Shops: {len(shops)}")
    for shop in shops:
        print(f"   • {shop['name']} ({shop['city']})")
    
    print(f"\n🥬 Groceries: {len(groceries)}")
    categories = {}
    for grocery in groceries:
        category = grocery['category']
        if category not in categories:
            categories[category] = []
        categories[category].append(grocery)
    
    for category, items in categories.items():
        organic_count = sum(1 for item in items if item['organic'])
        print(f"   • {category}: {len(items)} items ({organic_count} organic)")
    
    print(f"\n🛒 Shopping Events: {len(events)}")
    if events:
        total_spent = sum(event.get('total_amount', 0) for event in events)
        avg_spent = total_spent / len(events) if events else 0
        print(f"   • Total spent: ${total_spent:.2f}")
        print(f"   • Average per trip: ${avg_spent:.2f}")
        
        # Shop distribution
        shop_counts = {}
        for event in events:
            shop_name = event['shop']['name']
            shop_counts[shop_name] = shop_counts.get(shop_name, 0) + 1
        
        print("   • Events per shop:")
        for shop_name, count in sorted(shop_counts.items(), key=lambda x: x[1], reverse=True):
            print(f"     - {shop_name}: {count} events")

def main():
    """Main function to create all test data."""
    args = parse_arguments()
    
    print("🚀 GROCERY TRACKER TEST DATA GENERATOR")
    print("=" * 50)
    
    if args.dry_run:
        print("🔍 DRY RUN MODE - No data will be created")
    
    print(f"API URL: {args.url}")
    print(f"Shopping events: {args.events}")
    print(f"Date range: {args.days} days back")
    print()
    
    try:
        # Test connection
        if not check_api_connection(args.url):
            print(f"❌ Cannot connect to the API server at {args.url}")
            print("   Make sure the backend server is running!")
            sys.exit(1)
        
        print("✅ Connected to API server\n")
        
        shops = []
        groceries = []
        events = []
        
        # Create data based on arguments
        if args.shops_only:
            shops = create_shops(args.url, args.verbose, args.dry_run)
        elif args.groceries_only:
            groceries = create_groceries(args.url, args.verbose, args.dry_run)
        elif args.events_only:
            # Get existing data for events
            shops, groceries = get_existing_data(args.url)
            events = create_shopping_events(shops, groceries, args.url, args.events, args.days, args.verbose, args.dry_run)
        else:
            # Create all data
            shops = create_shops(args.url, args.verbose, args.dry_run)
            groceries = create_groceries(args.url, args.verbose, args.dry_run)
            if shops and groceries:
                events = create_shopping_events(shops, groceries, args.url, args.events, args.days, args.verbose, args.dry_run)
        
        # Print summary
        if not args.dry_run:
            print_summary(shops, groceries, events)
        
        if args.dry_run:
            print("\n🔍 Dry run completed. Use without --dry-run to actually create the data.")
        else:
            print("\n🎉 Test data creation completed successfully!")
            print("You can now explore the application with realistic data.")
        
    except KeyboardInterrupt:
        print("\n❌ Operation cancelled by user.")
        sys.exit(1)
    except requests.exceptions.ConnectionError:
        print(f"❌ Could not connect to the API server at {args.url}")
        print("   Make sure the backend server is running!")
        sys.exit(1)
    except Exception as e:
        print(f"❌ Unexpected error: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()