backend/assets/generate_and_save_coupons.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Script to generate and save coupons to the database
"""

import os
import sys
import random
from datetime import datetime, timedelta

# Add the backend directory to the path so we can import modules
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

# Import required modules
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.exc import SQLAlchemyError

from backend.app.admin.model.coupon import Coupon
from backend.utils.snowflake import snowflake
from backend.core.conf import settings, get_db_uri


def generate_coupon_codes(prefix: str, quantity: int):
    """
    Generate coupon codes with specified prefix and quantity.
    
    Format: [PREFIX][NUMBER] - Total 6 characters
    Example: A12345, TEST0, XYZ999
    
    Args:
        prefix (str): The letter prefix for the coupon codes (should be uppercase)
        quantity (int): Number of coupon codes to generate
        
    Returns:
        list: List of generated coupon codes
    """
    if not prefix.isalpha() or not prefix.isupper():
        raise ValueError("Prefix must be uppercase letters only")
    
    if len(prefix) == 0 or len(prefix) > 5:
        raise ValueError("Prefix must be 1-5 characters long")
    
    if quantity <= 0:
        raise ValueError("Quantity must be greater than 0")
    
    # Calculate number of digits based on prefix length (total 6 characters)
    num_digits = 6 - len(prefix)
    
    # Maximum possible combinations
    max_combinations = 10 ** num_digits
    
    if quantity > max_combinations:
        raise ValueError(f"With prefix '{prefix}' (length {len(prefix)}), can only generate {max_combinations} unique codes (0 to {max_combinations - 1})")
    
    codes = []
    # Generate incremental numbers starting from 0
    for i in range(quantity):
        # Format with leading zeros to make it the required number of digits
        formatted_number = f"{i:0{num_digits}d}"
        # Combine prefix with formatted number
        coupon_code = f"{prefix}{formatted_number}"
        codes.append(coupon_code)
    
    return codes


def save_coupons_to_db(prefix: str, quantity: int, coupon_type: str, points: int, expire_days: int = None):
    """
    Generate and save coupons to the database.
    
    Coupon codes are always 6 characters total:
    - 1-letter prefix: 5 digits (up to 100000 codes: A00000-A99999)
    - 4-letter prefix: 2 digits (up to 100 codes: TEST00-TEST99)
    - 5-letter prefix: 1 digit (up to 10 codes: ABCDE0-ABCDE9)
    
    Args:
        prefix (str): The letter prefix for the coupon codes
        quantity (int): Number of coupon codes to generate
        coupon_type (str): Type of the coupons
        points (int): Points value of the coupons
        expire_days (int, optional): Days until expiration. If None, no expiration.
    """
    # Create database engine and session
    db_url = get_db_uri(settings)
    # Replace asyncmy with mysql+mysqlconnector for synchronous connection
    sync_db_url = db_url.replace('mysql+asyncmy', 'mysql+mysqlconnector')
    
    try:
        engine = create_engine(sync_db_url, echo=False)
        SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
        db = SessionLocal()
        
        # Generate coupon codes
        codes = generate_coupon_codes(prefix, quantity)
        
        # Create coupon objects
        coupons = []
        for code in codes:
            # Generate snowflake ID
            coupon_id = snowflake.generate()
            
            # Calculate expiration date if needed
            expires_at = None
            if expire_days is not None and expire_days > 0:
                expires_at = datetime.now() + timedelta(days=expire_days)
            
            # Create coupon object
            # Note: id is auto-generated by snowflake, but we want to use our own snowflake generator
            coupon = Coupon(
                code=code,
                type=coupon_type,
                points=points,
                expires_at=expires_at
            )
            # Set the id manually after creation
            coupon.id = coupon_id
            coupons.append(coupon)
        
        # Bulk insert coupons
        db.add_all(coupons)
        db.commit()
        
        print(f"Successfully saved {len(coupons)} coupons to the database.")
        print(f"Prefix: {prefix}, Type: {coupon_type}, Points: {points}")
        if expire_days:
            print(f"Expires in: {expire_days} days")
        
        # Display first 5 coupons as examples
        print("\nSample coupons generated:")
        for coupon in coupons[:5]:
            print(f"  ID: {coupon.id}, Code: {coupon.code}")
        
        db.close()
        
    except SQLAlchemyError as e:
        print(f"Database error: {e}")
        if 'db' in locals():
            db.rollback()
            db.close()
    except Exception as e:
        print(f"Error: {e}")
        if 'db' in locals():
            db.close()


def main():
    """Main function to demonstrate usage"""
    print("Coupon Generator and Database Saver")
    print("=" * 40)
    
    # Example: Generate and save coupons with different prefixes
    try:
        # Single character prefix (5 digits, incremental from 00000)
        # print("Generating coupons with single character prefix 'A'...")
        # save_coupons_to_db('A', 5, 'NORMAL', 100, 30)
        # print("\n" + "-" * 40 + "\n")

        # 4-character prefix (2 digits, incremental from 00)
        print("Generating coupons with 4-character prefix 'TEST'...")
        save_coupons_to_db('VIP', 5, 'test', 1000, 60)
        print("\n" + "-" * 40 + "\n")
        
        # 3-character prefix (3 digits, incremental from 000)
        # print("Generating coupons with 3-character prefix 'XYZ'...")
        # save_coupons_to_db('XYZ', 3, 'SPECIAL', 500, 15)
        # print("\n" + "-" * 40 + "\n")
        
        # 5-character prefix (1 digit, incremental from 0)
        # print("Generating coupons with 5-character prefix 'ABCDE'...")
        # save_coupons_to_db('ABCDE', 5, 'PREMIUM', 2000, 90)
        
    except Exception as e:
        print(f"Error in main: {e}")


if __name__ == "__main__":
    main()