scripts/seed_data.py

"""
Seed database script to bootstrap the system with initial demo data.
Supports both PostgreSQL (local) and SQLite (Hugging Face / local) databases.
"""
from __future__ import annotations

import argparse
import asyncio
import sys
import os

from sqlalchemy import select

# Adjust path to import app modules
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "backend")))

from app.core.config import settings
from app.core.database import AsyncSessionLocal, engine
from app.core.security import UserRole, hash_password
from app.models.tenant import Tenant
from app.models.user import User


async def seed_database(hf_mode: bool = False):
    """Seed the database with a default tenant and administrator account."""
    print("Database seeding started...")
    
    # Override settings if hf_mode is specified
    if hf_mode:
        print("HF Mode enabled: seeding SQLite database...")
        
    async with AsyncSessionLocal() as session:
        try:
            # 1. Ensure default tenant exists
            tenant_slug = "acme"
            result = await session.execute(select(Tenant).where(Tenant.slug == tenant_slug))
            tenant = result.scalar_one_or_none()
            
            if not tenant:
                print("Creating default tenant 'Acme Corp'...")
                tenant = Tenant(
                    name="Acme Corp",
                    slug=tenant_slug,
                    description="Default demo organization workspace.",
                    is_active=True,
                    plan="free"
                )
                session.add(tenant)
                await session.flush()
            else:
                print(f"Default tenant 'Acme Corp' already exists (ID: {tenant.id})")

            # 2. Ensure default admin user exists
            admin_email = settings.FIRST_ADMIN_EMAIL or "admin@company.com"
            result = await session.execute(select(User).where(User.email == admin_email))
            admin_user = result.scalar_one_or_none()
            
            if not admin_user:
                print(f"Creating default admin user '{admin_email}'...")
                admin_password = settings.FIRST_ADMIN_PASSWORD or "Admin@123!"
                admin_user = User(
                    email=admin_email,
                    hashed_password=hash_password(admin_password),
                    full_name="System Admin",
                    role=UserRole.ADMIN.value,
                    tenant_id=tenant.id,
                    is_active=True,
                    is_verified=True
                )
                session.add(admin_user)
            else:
                print(f"Default admin user '{admin_email}' already exists.")

            await session.commit()
            print("Database seeding completed successfully! 🎉")
            
        except Exception as e:
            print(f"Error seeding database: {e}")
            await session.rollback()
            raise


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Seed database for Enterprise AI Copilot")
    parser.add_argument("--hf-mode", action="store_true", help="Run in Hugging Face SQLite mode")
    args = parser.parse_args()
    
    # Run async function
    asyncio.run(seed_database(hf_mode=args.hf_mode))