backend/db/seed.py

import os
import datetime
import random
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from backend.db.models import Base, DarkStore, Inventory, SalesEvent, Restaurant, Coupon, ExpenseLog, SystemSetting

DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://hyperflow_admin:hyperflow_secure_pass@localhost:5432/hyperflow_db")

def seed_database():
    print(f"Connecting to database at {DATABASE_URL}...")
    # Add connect_timeout to avoid blocking if postgres isn't running
    try:
        engine = create_engine(DATABASE_URL, connect_args={"connect_timeout": 5})
        # Check connection
        with engine.connect() as conn:
            pass
    except Exception as e:
        print(f"Could not connect to database: {e}. Skipping seed.")
        return

    # Ensure tables exist
    Base.metadata.create_all(bind=engine)
    
    Session = sessionmaker(bind=engine)
    session = Session()
    
    try:
        # Seed DarkStore records
        if session.query(DarkStore).first() is None:
            print("Seeding DarkStore records...")
            stores = [
                DarkStore(id="store_01", name="Whitefield Dark Store", city="Bengaluru", lat=12.9716, lng=77.5946),
                DarkStore(id="store_02", name="Koramangala Hub", city="Bengaluru", lat=12.9345, lng=77.6265),
                DarkStore(id="store_03", name="Indiranagar Dark Store", city="Bengaluru", lat=12.9784, lng=77.6408)
            ]
            session.add_all(stores)
            session.commit()
        
        if session.query(Inventory).first() is None:
            print("Seeding Inventory records...")
            # Add items. Make sure Amul Milk and Organic Bananas have 0 stock to trigger OOS
            inventory_items = [
                # store_01
                Inventory(store_id="store_01", sku_id="g1", sku_name="Fresh Toned Milk 1L", qty_available=0),
                Inventory(store_id="store_01", sku_id="g2", sku_name="Organic Bananas 1 Dozen", qty_available=0),
                Inventory(store_id="store_01", sku_id="g3", sku_name="Whole Wheat Bread 400g", qty_available=25),
                Inventory(store_id="store_01", sku_id="g4", sku_name="Spiced Chicken Burger Patty 4pcs", qty_available=15),
                
                # store_02
                Inventory(store_id="store_02", sku_id="g1", sku_name="Fresh Toned Milk 1L", qty_available=30),
                Inventory(store_id="store_02", sku_id="g2", sku_name="Organic Bananas 1 Dozen", qty_available=10),
                
                # store_03
                Inventory(store_id="store_03", sku_id="g1", sku_name="Fresh Toned Milk 1L", qty_available=50),
                Inventory(store_id="store_03", sku_id="g4", sku_name="Spiced Chicken Burger Patty 4pcs", qty_available=8)
            ]
            session.add_all(inventory_items)
            session.commit()
        
        if session.query(SalesEvent).count() < 50:
            print("Seeding 500+ SalesEvent historical records with weather & time features...")
            start_date = datetime.date.today() - datetime.timedelta(days=30)
            
            sales_events = []
            random.seed(42)
            for i in range(30):
                current_date = start_date + datetime.timedelta(days=i)
                for hour in [8, 10, 12, 14, 16, 18, 20]:
                    for store in ["store_01", "store_02", "store_03"]:
                        for sku in ["g1", "g2", "g3", "g4"]:
                            base_sales = 15.0 if sku in ["g1", "g2"] else 8.0
                            observed = max(0, float(random.normalvariate(base_sales, 4.0)))
                            
                            censored = False
                            oos_time = None
                            if sku in ["g1", "g2"] and random.random() < 0.35:
                                censored = True
                                observed = min(observed, 10.0)
                                oos_time = datetime.datetime.combine(current_date, datetime.time(hour, random.randint(10, 50)))
                            
                            temp = float(random.uniform(15.0, 38.0))
                            rain = float(random.exponential(2.0))
                            elapsed_sec = float(random.normalvariate(900.0, 300.0))
                            
                            sales_events.append(SalesEvent(
                                store_id=store,
                                sku_id=sku,
                                observed_sales=observed,
                                censored=censored,
                                oos_time=oos_time,
                                event_date=current_date,
                                hour_bucket=hour,
                                weather_temp=temp,
                                weather_rain=rain,
                                time_elapsed_sec=elapsed_sec
                            ))
            session.add_all(sales_events)
            session.commit()
            print(f"Successfully seeded {len(sales_events)} SalesEvent records for PSI monitoring!")

        # Seed Restaurants
        if session.query(Restaurant).first() is None:
            print("Seeding Restaurant records...")
            rests = [
                Restaurant(
                    id="rest_behrouz",
                    name="Behrouz Biryani",
                    cuisine="Biryani · Mughlai · Royal",
                    rating=4.6,
                    distance="2.1 km",
                    time="28 min",
                    slaConfidence=97,
                    isAIPick=True,
                    isExclusive=True,
                    image="https://lh3.googleusercontent.com/aida-public/AB6AXuB3O6h3kN5v2ZfZDd3Ufds1_PUUHBmlla4WShhsUOwN1BiWVty9aGs9k-ujSiY3HWg0c-a6yUVCpufZJTK3hqLopqOy-INM9HYG-SKcVE0PbA__mUudSLa2FZF4yeu1q6fwxpjVZXn7yNLyelP_KZmven-uKjmR8Q3bG2PkZi64JiSya_N0Zb1Ww0kf3A7LW34llf4b4dpiTff9GbejYkJFooJR4Slc4fs85sLnGz-kZjWnuFABxdtocK8oviRGW5vmkB6XF1IMU4YS"
                ),
                Restaurant(
                    id="rest_carbon_grill",
                    name="Carbon Grill",
                    cuisine="Burgers · Wings · Sides",
                    rating=4.3,
                    distance="1.4 km",
                    time="22 min",
                    slaConfidence=94,
                    isAIPick=False,
                    isExclusive=False,
                    image="https://lh3.googleusercontent.com/aida-public/AB6AXuD9C62CkwFO1Ta65rOPGt_zkQb3NWBfpIVfhSCWsS173P7Hw1t8O2CFnA1Swhsh03BFAJeCU4v8zMcs2FtgfS9UKrkQ-pgIxmQV0atKwEY1VvIrOO2nqjJirHB5LtlEy7v2E23zmpz5QUROCmGsEwpUTOxc6-W7bqEnwZTpjlEj84W0_wRNkm3oiChRsbQBbdUsj6iQ4IQ8MjgCXDjvXHjIGyb2EehurUmG2rcFE5E_2NQqMXhnC7sZPl5JUl0b-89s8s1A5HghkpjV"
                ),
                Restaurant(
                    id="rest_yoko_ono",
                    name="Yoko Ono Sushi",
                    cuisine="Sushi · Asian · Japanese",
                    rating=4.5,
                    distance="3.0 km",
                    time="32 min",
                    slaConfidence=96,
                    isAIPick=False,
                    isExclusive=True,
                    image="https://lh3.googleusercontent.com/aida-public/AB6AXuCBY63vuIkeBp6l5cHYDUYAUxyfZjekeIUDrgoaWXdYWfRsIItON9yVcNgasVY5EVJ_z9UCEYE7ifS6es_em8GXuQSZjL4elMAOcYKY-mFqvK7XoIYiCdoO9fXcs76s27BFjIlZ-jibt94sXMKAMiW-HDhL8Fx6YgFDMjXCKJuqgQvL6f2QokApfLDSvnpgf5uRCpVCyjlevWvENzKb2pD1gJvWBrOj_kU8HsHYg8siO1GP2yGFdEgOS79jFlelYdFjbEs_cIizY-X6"
                )
            ]
            session.add_all(rests)
            session.commit()

        # Seed Coupons
        if session.query(Coupon).first() is None:
            print("Seeding Coupon records...")
            coupons = [
                Coupon(code="SWIGGYIT", discount_percentage=50, min_cart_value=199.0, active=True),
                Coupon(code="JUMBO75", discount_percentage=75, min_cart_value=399.0, active=True)
            ]
            session.add_all(coupons)
            session.commit()

        # Seed Expense Logs
        if session.query(ExpenseLog).first() is None:
            print("Seeding ExpenseLog records...")
            expenses = [
                ExpenseLog(category="Food wastage claim", amount=2400.0, description="OOS threshold cleanup Whitefield Store"),
                ExpenseLog(category="Logistics rain incentive surge", amount=4120.0, description="Monsoon Storm Surge Fleet Payout")
            ]
            session.add_all(expenses)
            session.commit()

        # Seed System Settings
        if session.query(SystemSetting).first() is None:
            print("Seeding SystemSetting records...")
            settings = [
                SystemSetting(key="festival_theme", value="nominal")
            ]
            session.add_all(settings)
            session.commit()

        print("Database successfully seeded with historical operation parameters and admin baselines!")
    except Exception as e:
        session.rollback()
        print(f"Error during seeding: {e}")
        raise e
    finally:
        session.close()

if __name__ == "__main__":
    seed_database()