thrillwiki_django_no_react/parks/management/commands/seed_data.py

import os
import json
import random
from datetime import datetime
from django.core.management.base import BaseCommand
from django.contrib.auth import get_user_model
from django.core.files import File
from django.utils.text import slugify
from django.contrib.contenttypes.models import ContentType
from faker import Faker
import requests
from io import BytesIO
from PIL import Image

from parks.models import Park
from rides.models import Ride, RollerCoasterStats
from reviews.models import Review
from media.models import Photo
from accounts.models import User, UserProfile, TopList, TopListItem
from companies.models import Company, Manufacturer

fake = Faker()

class Command(BaseCommand):
    help = 'Seeds the database with sample data'

    def add_arguments(self, parser):
        parser.add_argument('--users', type=int, default=50)
        parser.add_argument('--reviews-per-item', type=int, default=10)

    def download_and_save_image(self, url, prefix):
        try:
            response = requests.get(url)
            img = Image.open(BytesIO(response.content))
            img_io = BytesIO()
            img.save(img_io, format='JPEG')
            img_io.seek(0)
            return f'{prefix}_{fake.uuid4()}.jpg', File(img_io)
        except:
            return None, None

    def create_users(self, count):
        self.stdout.write('Creating users...')
        users = []
        roles = ['USER'] * 20 + ['MODERATOR'] * 3 + ['ADMIN'] * 2
        
        # Create or get superuser
        try:
            superuser = User.objects.get(username='admin')
            self.stdout.write('Superuser already exists')
        except User.DoesNotExist:
            superuser = User.objects.create_superuser(
                username='admin',
                email='admin@thrillwiki.com',
                password='admin',
                role='SUPERUSER'
            )
            UserProfile.objects.create(
                user=superuser,
                display_name='Admin',
                pronouns='they/them',
                bio='ThrillWiki Administrator'
            )
            self.stdout.write('Created superuser')
        
        users.append(superuser)

        # Delete existing non-superuser users if any
        User.objects.exclude(username='admin').delete()
        self.stdout.write('Deleted existing users')

        for _ in range(count):
            username = fake.user_name()
            while User.objects.filter(username=username).exists():
                username = fake.user_name()
            
            user = User.objects.create_user(
                username=username,
                email=fake.email(),
                password='password123',
                role=random.choice(roles)
            )
            
            # Create user profile
            profile = UserProfile.objects.create(
                user=user,
                display_name=fake.name(),
                pronouns=random.choice(['he/him', 'she/her', 'they/them', '']),
                bio=fake.text(max_nb_chars=200),
                twitter=fake.url() if random.choice([True, False]) else '',
                instagram=fake.url() if random.choice([True, False]) else '',
                youtube=fake.url() if random.choice([True, False]) else '',
                discord=fake.user_name() if random.choice([True, False]) else '',
                coaster_credits=random.randint(0, 500),
                dark_ride_credits=random.randint(0, 200),
                flat_ride_credits=random.randint(0, 300),
                water_ride_credits=random.randint(0, 100)
            )
            
            # Add avatar
            img_url = f'https://picsum.photos/200/200?random={fake.random_number(5)}'
            filename, file = self.download_and_save_image(img_url, 'avatar')
            if filename and file:
                profile.avatar.save(filename, file, save=True)
            
            users.append(user)
            self.stdout.write(f'Created user: {username}')
        
        return users

    def create_companies(self):
        self.stdout.write('Creating companies...')
        # Delete existing companies
        Company.objects.all().delete()
        self.stdout.write('Deleted existing companies')

        companies = {
            'The Walt Disney Company': {
                'headquarters': 'Burbank, California',
                'founded_date': '1923-10-16',
                'website': 'https://www.disney.com',
            },
            'Cedar Fair': {
                'headquarters': 'Sandusky, Ohio',
                'founded_date': '1983-05-01',
                'website': 'https://www.cedarfair.com',
            },
            'NBCUniversal': {
                'headquarters': 'New York City, New York',
                'founded_date': '1912-04-30',
                'website': 'https://www.nbcuniversal.com',
            },
            'Merlin Entertainments': {
                'headquarters': 'Poole, England',
                'founded_date': '1999-05-19',
                'website': 'https://www.merlinentertainments.biz',
            },
            'Mack Rides': {
                'headquarters': 'Waldkirch, Germany',
                'founded_date': '1780-01-01',
                'website': 'https://mack-rides.com',
            },
        }

        company_instances = {}
        for name, details in companies.items():
            company = Company.objects.create(
                name=name,
                slug=slugify(name),
                headquarters=details['headquarters'],
                founded_date=datetime.strptime(details['founded_date'], '%Y-%m-%d').date(),
                website=details['website'],
            )
            company_instances[name] = company
            self.stdout.write(f'Created company: {name}')

        return company_instances

    def create_manufacturers(self):
        self.stdout.write('Creating manufacturers...')
        # Delete existing manufacturers
        Manufacturer.objects.all().delete()
        self.stdout.write('Deleted existing manufacturers')

        manufacturers = {
            'Walt Disney Imagineering': {
                'headquarters': 'Glendale, California',
                'founded_date': '1952-12-16',
                'website': 'https://sites.disney.com/waltdisneyimagineering/',
            },
            'Bolliger & Mabillard': {
                'headquarters': 'Monthey, Switzerland',
                'founded_date': '1988-01-01',
                'website': 'https://www.bolliger-mabillard.com',
            },
            'Intamin': {
                'headquarters': 'Schaan, Liechtenstein',
                'founded_date': '1967-01-01',
                'website': 'https://www.intamin.com',
            },
            'Rocky Mountain Construction': {
                'headquarters': 'Hayden, Idaho',
                'founded_date': '2001-01-01',
                'website': 'https://www.rockymountainconstruction.com',
            },
            'Vekoma': {
                'headquarters': 'Vlodrop, Netherlands',
                'founded_date': '1926-01-01',
                'website': 'https://www.vekoma.com',
            },
            'Mack Rides': {
                'headquarters': 'Waldkirch, Germany',
                'founded_date': '1780-01-01',
                'website': 'https://mack-rides.com',
            },
            'Oceaneering International': {
                'headquarters': 'Houston, Texas',
                'founded_date': '1964-01-01',
                'website': 'https://www.oceaneering.com',
            },
        }

        manufacturer_instances = {}
        for name, details in manufacturers.items():
            manufacturer = Manufacturer.objects.create(
                name=name,
                slug=slugify(name),
                headquarters=details['headquarters'],
                founded_date=datetime.strptime(details['founded_date'], '%Y-%m-%d').date(),
                website=details['website'],
            )
            manufacturer_instances[name] = manufacturer
            self.stdout.write(f'Created manufacturer: {name}')

        return manufacturer_instances

    def create_parks_and_rides(self, users):
        self.stdout.write('Creating parks and rides from seed data...')
        
        # Create companies and manufacturers first
        companies = self.create_companies()
        manufacturers = self.create_manufacturers()
        
        # Load seed data
        seed_data_path = os.path.join(os.path.dirname(__file__), 'seed_data.json')
        with open(seed_data_path, 'r') as f:
            seed_data = json.load(f)
        
        # Delete existing parks (this will cascade delete rides)
        Park.objects.all().delete()
        self.stdout.write('Deleted existing parks and rides')
        
        parks = []
        for park_data in seed_data['parks']:
            # Create park with company instance
            park = Park.objects.create(
                name=park_data['name'],
                slug=slugify(park_data['name']),
                location=park_data['location'],
                country=park_data['country'],
                opening_date=datetime.strptime(park_data['opening_date'], '%Y-%m-%d').date(),
                status=park_data['status'],
                description=park_data['description'],
                website=park_data['website'],
                owner=companies[park_data['owner']],
                size_acres=park_data['size_acres']
            )
            
            # Add park photos
            for _ in range(random.randint(2, 5)):
                img_url = f'https://picsum.photos/800/600?random={fake.random_number(5)}'
                filename, file = self.download_and_save_image(img_url, 'park')
                if filename and file:
                    Photo.objects.create(
                        content_object=park,
                        image=file,
                        uploaded_by=random.choice(users),
                        caption=fake.sentence(),
                        is_approved=True
                    )
            
            # Create rides for this park
            for ride_data in park_data['rides']:
                ride = Ride.objects.create(
                    name=ride_data['name'],
                    slug=slugify(ride_data['name']),
                    category=ride_data['category'],
                    park=park,
                    status=ride_data['status'],
                    opening_date=datetime.strptime(ride_data['opening_date'], '%Y-%m-%d').date(),
                    manufacturer=manufacturers[ride_data['manufacturer']],
                    description=ride_data['description']
                )
                
                # Add roller coaster stats if applicable
                if ride_data['category'] == 'RC' and 'stats' in ride_data:
                    RollerCoasterStats.objects.create(
                        ride=ride,
                        height_ft=ride_data['stats']['height_ft'],
                        length_ft=ride_data['stats']['length_ft'],
                        speed_mph=ride_data['stats']['speed_mph'],
                        inversions=ride_data['stats']['inversions'],
                        ride_time_seconds=ride_data['stats']['ride_time_seconds']
                    )
                
                # Add ride photos
                for _ in range(random.randint(2, 5)):
                    img_url = f'https://picsum.photos/800/600?random={fake.random_number(5)}'
                    filename, file = self.download_and_save_image(img_url, 'ride')
                    if filename and file:
                        Photo.objects.create(
                            content_object=ride,
                            image=file,
                            uploaded_by=random.choice(users),
                            caption=fake.sentence(),
                            is_approved=True
                        )
            
            parks.append(park)
            self.stdout.write(f'Created park and rides: {park.name}')
        
        return parks

    def create_reviews(self, users, reviews_per_item):
        self.stdout.write('Creating reviews...')
        # Delete existing reviews
        Review.objects.all().delete()
        self.stdout.write('Deleted existing reviews')
        
        # Park reviews
        total_parks = Park.objects.count()
        for i, park in enumerate(Park.objects.all(), 1):
            for _ in range(random.randint(reviews_per_item - 5, reviews_per_item + 5)):
                Review.objects.create(
                    user=random.choice(users),
                    content_object=park,
                    title=fake.sentence(),
                    content=fake.text(max_nb_chars=500),
                    rating=random.randint(1, 10),
                    visit_date=fake.date_between(start_date=park.opening_date, end_date='today'),
                    is_published=True
                )
            if i % 5 == 0:
                self.stdout.write(f'Created reviews for {i}/{total_parks} parks')

        # Ride reviews
        total_rides = Ride.objects.count()
        for i, ride in enumerate(Ride.objects.all(), 1):
            for _ in range(random.randint(reviews_per_item - 5, reviews_per_item + 5)):
                Review.objects.create(
                    user=random.choice(users),
                    content_object=ride,
                    title=fake.sentence(),
                    content=fake.text(max_nb_chars=500),
                    rating=random.randint(1, 10),
                    visit_date=fake.date_between(start_date=ride.opening_date, end_date='today'),
                    is_published=True
                )
            if i % 20 == 0:
                self.stdout.write(f'Created reviews for {i}/{total_rides} rides')

    def create_top_lists(self, users):
        self.stdout.write('Creating top lists...')
        # Delete existing top lists
        TopList.objects.all().delete()
        self.stdout.write('Deleted existing top lists')
        
        categories = ['RC', 'DR', 'FR', 'WR', 'PK']
        total_users = len(users)
        
        # Get content types
        park_ct = ContentType.objects.get_for_model(Park)
        ride_ct = ContentType.objects.get_for_model(Ride)
        
        for i, user in enumerate(users, 1):
            for category in categories:
                if random.choice([True, False]):  # 50% chance to create a list
                    top_list = TopList.objects.create(
                        user=user,
                        title=f"My Top {random.randint(5, 20)} {dict(TopList.Categories.choices)[category]}s",
                        category=category,
                        description=fake.text(max_nb_chars=200)
                    )
                    
                    # Add items to the list
                    items = []
                    if category == 'PK':
                        items = list(Park.objects.all())
                        content_type = park_ct
                    else:
                        items = list(Ride.objects.filter(category=category))
                        content_type = ride_ct
                    
                    if items:
                        selected_items = random.sample(items, min(len(items), random.randint(5, 20)))
                        for rank, item in enumerate(selected_items, 1):
                            TopListItem.objects.create(
                                top_list=top_list,
                                content_type=content_type,
                                object_id=item.id,
                                rank=rank,
                                notes=fake.sentence() if random.choice([True, False]) else ''
                            )
            
            if i % 10 == 0:
                self.stdout.write(f'Created top lists for {i}/{total_users} users')

    def handle(self, *args, **options):
        self.stdout.write('Starting database seed...')
        
        users = self.create_users(options['users'])
        parks = self.create_parks_and_rides(users)
        self.create_reviews(users, options['reviews_per_item'])
        self.create_top_lists(users)
        
        self.stdout.write(self.style.SUCCESS('Successfully seeded database'))