seed_data.py

"""
Seed database with initial data for testing
"""
import sys
from pathlib import Path

# Add parent directory to path
sys.path.append(str(Path(__file__).parent))

from app.database import SessionLocal, PlantLocation, RateLimitConfig
from app.middleware import add_whitelist_ip
from app.auth import get_password_hash
from app.database import User

def seed_database():
    """Seed database with test data"""
    db = SessionLocal()
    
    try:
        print("Seeding database...")
        
        # Seed plant locations
        if not db.query(PlantLocation).filter(PlantLocation.plant_id == "plant_a").first():
            plant_a = PlantLocation(
                plant_id="plant_a",
                country_code="IN",
                country_name="India",
                city="Mumbai",
                description="Manufacturing Plant A - Mumbai, Maharashtra"
            )
            db.add(plant_a)
            print("✅ Added Plant A location")
        
        if not db.query(PlantLocation).filter(PlantLocation.plant_id == "plant_b").first():
            plant_b = PlantLocation(
                plant_id="plant_b",
                country_code="IN",
                country_name="India",
                city="Delhi",
                description="Manufacturing Plant B - Delhi, NCR"
            )
            db.add(plant_b)
            print("✅ Added Plant B location")
        
        if not db.query(PlantLocation).filter(PlantLocation.plant_id == "plant_c").first():
            plant_c = PlantLocation(
                plant_id="plant_c",
                country_code="IN",
                country_name="India",
                city="Bangalore",
                description="Manufacturing Plant C - Bangalore, Karnataka"
            )
            db.add(plant_c)
            print("✅ Added Plant C location")
        
        # Seed rate limit configs
        if not db.query(RateLimitConfig).filter(RateLimitConfig.endpoint == "/api/v1/sync").first():
            sync_config = RateLimitConfig(
                endpoint="/api/v1/sync",
                requests_per_minute=100,
                priority="critical",
                enabled=True
            )
            db.add(sync_config)
            print("✅ Added rate limit config for /api/v1/sync")
        
        if not db.query(RateLimitConfig).filter(RateLimitConfig.endpoint == "/token").first():
            auth_config = RateLimitConfig(
                endpoint="/token",
                requests_per_minute=10,
                priority="normal",
                enabled=True
            )
            db.add(auth_config)
            print("✅ Added rate limit config for /token")
        
        # Seed additional users
        if not db.query(User).filter(User.username == "plant_b").first():
            plant_b_user = User(
                username="plant_b",
                password_hash=get_password_hash("plant123"),
                role="plant",
                plant_id="plant_b"
            )
            db.add(plant_b_user)
            print("✅ Added Plant B user")
        
        if not db.query(User).filter(User.username == "plant_c").first():
            plant_c_user = User(
                username="plant_c",
                password_hash=get_password_hash("plant123"),
                role="plant",
                plant_id="plant_c"
            )
            db.add(plant_c_user)
            print("✅ Added Plant C user")
        
        db.commit()
        
        # Seed whitelisted IPs (using middleware function)
        print("\nSeeding whitelisted IPs...")
        add_whitelist_ip("127.0.0.1", added_by="system", notes="Localhost for development")
        print("✅ Added 127.0.0.1 to whitelist")
        
        print("\n✅ Database seeding completed successfully!")
        
    except Exception as e:
        print(f"\n❌ Error seeding database: {e}")
        db.rollback()
        raise
    finally:
        db.close()

if __name__ == "__main__":
    seed_database()