WealthWise/backend/scripts/seed_db.py

"""Database seeding script for WealthWise.

This script populates the database with sample data for development and testing.
"""

import asyncio
from datetime import date, timedelta
from decimal import Decimal
from uuid import uuid4

from sqlalchemy import select
from sqlmodel.ext.asyncio.session import AsyncSession

from app.core.db import AsyncSessionLocal, engine
from app.core.security import get_password_hash
from app.models import Portfolio, Transaction, User
from app.models.transaction import TransactionType


async def seed_database():
    """Seed the database with sample data."""
    print("Starting database seeding...")
    
    async with AsyncSessionLocal() as session:
        # Check if demo user already exists
        result = await session.execute(
            select(User).where(User.email == "demo@wealthwise.app")
        )
        existing_user = result.scalar_one_or_none()
        
        if existing_user:
            print("Demo user already exists. Skipping seeding.")
            return
        
        # Create demo user
        print("Creating demo user...")
        demo_user = User(
            id=uuid4(),
            email="demo@wealthwise.app",
            hashed_password=get_password_hash("password123"),
            is_active=True,
            is_superuser=False,
        )
        session.add(demo_user)
        await session.flush()  # Flush to get the user ID
        print(f"Created user: {demo_user.email} (ID: {demo_user.id})")
        
        # Create portfolio
        print("Creating portfolio...")
        portfolio = Portfolio(
            id=uuid4(),
            user_id=demo_user.id,
            name="Main Investment",
            description="Primary investment portfolio for tracking stocks, mutual funds, and other assets.",
        )
        session.add(portfolio)
        await session.flush()
        print(f"Created portfolio: {portfolio.name} (ID: {portfolio.id})")
        
        # Create sample transactions
        print("Creating transactions...")
        transactions = [
            # Income transactions
            Transaction(
                id=uuid4(),
                portfolio_id=portfolio.id,
                amount=Decimal("75000.00"),
                transaction_type=TransactionType.INCOME,
                transaction_date=date.today() - timedelta(days=30),
                description="Monthly Salary",
                category="Salary",
            ),
            Transaction(
                id=uuid4(),
                portfolio_id=portfolio.id,
                amount=Decimal("2500.00"),
                transaction_type=TransactionType.INCOME,
                transaction_date=date.today() - timedelta(days=25),
                description="Freelance Project Payment",
                category="Freelance",
            ),
            Transaction(
                id=uuid4(),
                portfolio_id=portfolio.id,
                amount=Decimal("1200.00"),
                transaction_type=TransactionType.INCOME,
                transaction_date=date.today() - timedelta(days=20),
                description="Stock Dividend - TCS",
                category="Dividend",
            ),
            # Expense transactions
            Transaction(
                id=uuid4(),
                portfolio_id=portfolio.id,
                amount=Decimal("25000.00"),
                transaction_type=TransactionType.EXPENSE,
                transaction_date=date.today() - timedelta(days=28),
                description="Monthly Rent",
                category="Housing",
            ),
            Transaction(
                id=uuid4(),
                portfolio_id=portfolio.id,
                amount=Decimal("8000.00"),
                transaction_type=TransactionType.EXPENSE,
                transaction_date=date.today() - timedelta(days=26),
                description="Grocery Shopping - BigBasket",
                category="Groceries",
            ),
            Transaction(
                id=uuid4(),
                portfolio_id=portfolio.id,
                amount=Decimal("3500.00"),
                transaction_type=TransactionType.EXPENSE,
                transaction_date=date.today() - timedelta(days=22),
                description="Electricity & Internet Bill",
                category="Utilities",
            ),
            Transaction(
                id=uuid4(),
                portfolio_id=portfolio.id,
                amount=Decimal("5000.00"),
                transaction_type=TransactionType.EXPENSE,
                transaction_date=date.today() - timedelta(days=18),
                description="Dining Out & Entertainment",
                category="Entertainment",
            ),
            # Investment transactions
            Transaction(
                id=uuid4(),
                portfolio_id=portfolio.id,
                amount=Decimal("10000.00"),
                transaction_type=TransactionType.TRANSFER,
                transaction_date=date.today() - timedelta(days=24),
                description="Stock Purchase - Infosys",
                category="Investment",
            ),
            Transaction(
                id=uuid4(),
                portfolio_id=portfolio.id,
                amount=Decimal("5000.00"),
                transaction_type=TransactionType.TRANSFER,
                transaction_date=date.today() - timedelta(days=21),
                description="Mutual Fund SIP - SBI Bluechip",
                category="Investment",
            ),
            Transaction(
                id=uuid4(),
                portfolio_id=portfolio.id,
                amount=Decimal("3000.00"),
                transaction_type=TransactionType.TRANSFER,
                transaction_date=date.today() - timedelta(days=15),
                description="Crypto Purchase - Bitcoin",
                category="Investment",
            ),
        ]
        
        for transaction in transactions:
            session.add(transaction)
        
        await session.commit()
        print(f"Created {len(transactions)} transactions")
        
        # Calculate summary
        income_total = sum(t.amount for t in transactions if t.transaction_type == TransactionType.INCOME)
        expense_total = sum(t.amount for t in transactions if t.transaction_type == TransactionType.EXPENSE)
        transfer_total = sum(t.amount for t in transactions if t.transaction_type == TransactionType.TRANSFER)
        
        print("\n" + "="*50)
        print("SEEDING COMPLETE!")
        print("="*50)
        print(f"User: {demo_user.email}")
        print(f"Password: password123")
        print(f"Portfolio: {portfolio.name}")
        print(f"\nTransaction Summary:")
        print(f"  Total Income: ₹{income_total:,.2f}")
        print(f"  Total Expenses: ₹{expense_total:,.2f}")
        print(f"  Total Investments: ₹{transfer_total:,.2f}")
        print(f"  Net Cash Flow: ₹{(income_total - expense_total):,.2f}")
        print("="*50)


async def main():
    """Main entry point for seeding."""
    try:
        await seed_database()
    except Exception as e:
        print(f"Error seeding database: {e}")
        raise
    finally:
        await engine.dispose()


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