gen.py

#!/usr/bin/env python3
"""
gen.py

Async script to generate test data for orders service using asyncpg.
- Products and Product Categories
- Orders and Order Items

Assumes all tables already exist in the database.
Displays console progress bars using tqdm.

Usage:
  python gen.py --num-products 50000 --num-orders 500000 --avg-items-per-order 2
"""

import argparse
import asyncio
import math
import random
import uuid
from datetime import datetime, timedelta

import asyncpg
from tqdm import tqdm

from orders.store.config import DATABASE_URL as ORDERS_DATABASE_URL


def parse_args():
    parser = argparse.ArgumentParser(
        description="Generate test data for orders and order items asynchronously."
    )
    parser.add_argument(
        "--num-products",
        type=int,
        default=50000,
        help="Number of products to generate",
    )
    parser.add_argument(
        "--num-categories",
        type=int,
        default=20,
        help="Number of product categories to generate",
    )
    parser.add_argument(
        "--num-orders",
        type=int,
        default=10000,
        help="Number of orders to generate",
    )
    parser.add_argument(
        "--avg-items-per-order",
        type=float,
        default=3.0,
        help="Average number of items per order",
    )
    parser.add_argument(
        "--num-users",
        type=int,
        default=None,
        help="Number of unique users to generate (default: max(100, num_orders // 10))",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=100_000,
        help="Batch insert size for all tables",
    )
    return parser.parse_args()


def poisson(lmbda: float) -> int:
    """
    Knuth's algorithm for Poisson distribution.
    """
    L = math.exp(-lmbda)
    k = 0
    p = 1.0
    while p > L:
        k += 1
        p *= random.random()
    return max(0, k - 1)


async def generate_categories_data(conn, args):
    """Generate product categories."""
    print(f"Generating {args.num_categories} product categories...")
    categories = [
        "Electronics",
        "Clothing",
        "Books",
        "Home & Garden",
        "Sports",
        "Toys",
        "Food & Beverages",
        "Beauty",
        "Automotive",
        "Health",
        "Office",
        "Pet Supplies",
        "Jewelry",
        "Music",
        "Movies",
        "Games",
        "Tools",
        "Baby",
        "Outdoor",
        "Art & Crafts",
    ]

    categories_batch = []
    for i in range(1, args.num_categories + 1):
        name = categories[i % len(categories)] + (
            f" {i // len(categories)}" if i > len(categories) else ""
        )
        categories_batch.append((i, name))

    await conn.executemany(
        "INSERT INTO products_categories(id, name) VALUES($1, $2)",
        categories_batch,
    )
    print(f"Inserted {len(categories_batch)} categories.")


async def generate_products_data(conn, args):
    """Generate products."""
    print(f"Generating {args.num_products} products...")
    products_batch = []

    for i in tqdm(range(1, args.num_products + 1), desc="Products", unit="product"):
        title = f"Product {i}"
        price = random.randint(100, 50000)  # Price in cents (1-500 rubles)
        photo_url = f"https://example.com/photos/product_{i}.jpg"
        category_id = random.randint(1, args.num_categories)

        products_batch.append((i, title, price, photo_url, category_id))

        if len(products_batch) >= args.batch_size:
            await conn.executemany(
                "INSERT INTO products(id, title, price, photo_url, category_id) VALUES($1, $2, $3, $4, $5)",
                products_batch,
            )
            products_batch.clear()

    if products_batch:
        await conn.executemany(
            "INSERT INTO products(id, title, price, photo_url, category_id) VALUES($1, $2, $3, $4, $5)",
            products_batch,
        )

    print(f"Inserted {args.num_products} products.")


async def generate_orders_data(conn, args):
    """Generate orders and order items data for orders service."""
    # Generate orders with progress bar
    print(
        f"Generating {args.num_orders} orders (avg {args.avg_items_per_order} items each)..."
    )
    orders_batch = []
    order_items_batch = []
    total_order_items = 0
    now = datetime.now()

    # Generate unique user IDs
    num_users = (
        args.num_users
        if args.num_users is not None
        else max(100, args.num_orders // 10)
    )
    user_ids = [uuid.uuid4() for _ in range(num_users)]
    print(f"Using {num_users} unique users")

    # Product ID range for random selection
    min_product_id = 1
    max_product_id = args.num_products

    for _ in tqdm(range(args.num_orders), desc="Orders", unit="order"):
        order_id = uuid.uuid4()
        user_id = random.choice(user_ids)
        created = now - timedelta(days=random.randint(0, 365))

        # Generate order items for this order first to calculate total price
        num_items = poisson(args.avg_items_per_order)
        num_items = max(1, num_items)

        order_total_price = 0
        used_pids = set()

        for _ in range(num_items):
            # Быстрый выбор случайного товара
            pid = random.randint(min_product_id, max_product_id)
            if pid in used_pids:
                continue
            used_pids.add(pid)

            qty = random.randint(1, 5)
            item_price = random.randint(100, 5000)  # Price in cents (1-50 rubles)
            order_items_batch.append((str(order_id), pid, qty, item_price))
            order_total_price += qty * item_price
            total_order_items += 1

        # Add order with calculated total price
        orders_batch.append((str(order_id), str(user_id), created, order_total_price))

        # Batch insert order items
        if len(order_items_batch) >= args.batch_size:
            if orders_batch:
                await conn.executemany(
                    "INSERT INTO orders(id, user_id, created_at, total_price) VALUES($1, $2, $3, $4)",
                    orders_batch,
                )
                orders_batch.clear()
            await conn.executemany(
                "INSERT INTO order_items(order_id, product_id, quantity, price) VALUES($1, $2, $3, $4)",
                order_items_batch,
            )
            order_items_batch.clear()

    # Insert remaining batches
    if orders_batch:
        await conn.executemany(
            "INSERT INTO orders(id, user_id, created_at, total_price) VALUES($1, $2, $3, $4)",
            orders_batch,
        )

    if order_items_batch:
        await conn.executemany(
            "INSERT INTO order_items(order_id, product_id, quantity, price) VALUES($1, $2, $3, $4)",
            order_items_batch,
        )

    print(f"Inserted {args.num_orders} orders and {total_order_items} order items.")


async def main():
    args = parse_args()

    print("=== Generating Test Data ===")
    # Connect to orders database
    orders_conn = await asyncpg.connect(
        dsn=ORDERS_DATABASE_URL.replace("postgresql+asyncpg", "postgresql")
    )

    # Generate data in order: categories -> products -> orders -> order_items
    await generate_categories_data(orders_conn, args)
    await generate_products_data(orders_conn, args)
    await generate_orders_data(orders_conn, args)

    await orders_conn.close()

    print("\n=== Data Generation Complete ===")


if __name__ == "__main__":
    asyncio.run(main())