seed_plans.py

"""
Seed script to create default subscription plans.
Run with: poetry run python seed_plans.py
"""
from app.db.session import SessionLocal
# Import all models to ensure relationships are resolved
from app.models.base import *  # This imports all models
from app.models.subscription_plan import SubscriptionPlan


def seed_plans():
    """Create default subscription plans if they don't exist."""
    db = SessionLocal()
    
    try:
        # Check if plans already exist
        existing_free = db.query(SubscriptionPlan).filter(SubscriptionPlan.name == "free").first()
        existing_test_pro = db.query(SubscriptionPlan).filter(SubscriptionPlan.name == "test pro").first()
        existing_pro = db.query(SubscriptionPlan).filter(SubscriptionPlan.name == "pro").first()
        existing_enterprise = db.query(SubscriptionPlan).filter(SubscriptionPlan.name == "enterprise").first()
        
        if not existing_free:
            free_plan = SubscriptionPlan(
                name="free",
                display_name="Free",
                price_monthly=0,
                price_yearly=0,
                max_conversations=3,
                max_storage_bytes=5242880,  # 5MB
                max_generated_files=3,
                features={
                    # "support": "community",
                    # "operations": ["split", "merge", "compress", "add_signature"]
                },
                is_active=True
            )
            db.add(free_plan)
            print("✅ Created Free plan")
        else:
            print("ℹ️ Free plan already exists")
        
        if not existing_test_pro:
            test_pro_plan = SubscriptionPlan(
                name="test pro",
                display_name="Test Pro",
                razorpay_plan_id=None,  # Set this after creating plan in Razorpay Dashboard
                price_monthly=499,  # ₹499
                price_yearly=4999,  # ₹4999
                max_conversations=-1,  # Unlimited
                max_storage_bytes=-1,  # Unlimited
                max_generated_files=-1,  # Unlimited
                features={
                    # "support": "priority",
                    # "operations": ["split", "merge", "compress", "add_signature"],
                    # "priority_processing": True
                },
                is_active=True
            )
            db.add(test_pro_plan)
            print("✅ Created Test Pro plan")
            print("⚠️  Remember to set razorpay_plan_id after creating the plan in Razorpay Dashboard!")
        else:
            print("ℹ️ Test Pro plan already exists")
        
        if not existing_pro:
            pro_plan = SubscriptionPlan(
                name="pro",
                display_name="Pro",
                razorpay_plan_id="plan_S4yvjZ4rynMQqV",
                price_monthly=499,  # ₹499
                price_yearly=4999,  # ₹4999
                max_conversations=-1,  # Unlimited
                max_storage_bytes=-1,  # Unlimited
                max_generated_files=-1,  # Unlimited
                features={
                    # "support": "priority",
                    # "operations": ["split", "merge", "compress", "add_signature"],
                    # "priority_processing": True
                },
                is_active=True
            )
            db.add(pro_plan)
            print("✅ Created Pro plan with razorpay_plan_id: plan_S4yvjZ4rynMQqV")
        else:
            print("ℹ️ Pro plan already exists")
        
        if not existing_enterprise:
            enterprise_plan = SubscriptionPlan(
                name="enterprise",
                display_name="Enterprise",
                price_monthly=-1,  # Custom pricing / Contact sales
                price_yearly=-1,
                max_conversations=-1,  # Unlimited
                max_storage_bytes=-1,  # Unlimited
                max_generated_files=-1,  # Unlimited
                features={
                    # "support": "dedicated",
                    # "api_access": True,
                    # "custom_branding": True,
                    # "sso": True
                },
                is_active=True
            )
            db.add(enterprise_plan)
            print("✅ Created Enterprise plan")
        else:
            print("ℹ️ Enterprise plan already exists")
        
        db.commit()
        print("\n🎉 Seed completed successfully!")
        
        # Display all plans
        print("\n📋 Current plans:")
        plans = db.query(SubscriptionPlan).all()
        for plan in plans:
            print(f"  - {plan.display_name} ({plan.name})")
            print(f"    Max conversations: {plan.max_conversations if plan.max_conversations != -1 else 'Unlimited'}")
            print(f"    Max storage: {plan.max_storage_bytes if plan.max_storage_bytes != -1 else 'Unlimited'}")
            print(f"    Max generated files: {plan.max_generated_files if plan.max_generated_files != -1 else 'Unlimited'}")
            print(f"    Price: ₹{plan.price_monthly}/month")
            print()
        
    except Exception as e:
        print(f"❌ Error: {e}")
        db.rollback()
    finally:
        db.close()


def update_razorpay_plan_id(plan_name: str, razorpay_plan_id: str):
    """Update a plan's razorpay_plan_id after creating it in Razorpay Dashboard."""
    db = SessionLocal()
    try:
        plan = db.query(SubscriptionPlan).filter(SubscriptionPlan.name == plan_name).first()
        if not plan:
            print(f"❌ Plan '{plan_name}' not found")
            return
        
        plan.razorpay_plan_id = razorpay_plan_id
        db.commit()
        print(f"✅ Updated {plan_name} with razorpay_plan_id: {razorpay_plan_id}")
    except Exception as e:
        print(f"❌ Error: {e}")
        db.rollback()
    finally:
        db.close()


if __name__ == "__main__":
    import sys
    
    if len(sys.argv) == 3:
        # Update razorpay_plan_id: poetry run python seed_plans.py pro plan_xxxxx
        plan_name = sys.argv[1]
        razorpay_plan_id = sys.argv[2]
        update_razorpay_plan_id(plan_name, razorpay_plan_id)
    else:
        seed_plans()
        print("\n💡 To update razorpay_plan_id, run:")
        print("   poetry run python seed_plans.py pro plan_xxxxx")