database_setup.py

import asyncio
import os
import logging
import uuid
from datetime import datetime, timedelta
from motor.motor_asyncio import AsyncIOMotorClient
from pymongo import IndexModel, ASCENDING, DESCENDING, TEXT
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Configure Logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

MONGODB_URL = os.getenv("MONGODB_URL", "mongodb://localhost:27017")
DATABASE_NAME = "nadraguard_db"

async def setup_database():
    logger.info("Initializing NadraGuard Database Setup...")
    client = AsyncIOMotorClient(MONGODB_URL)
    db = client[DATABASE_NAME]

    # --- 1. USERS COLLECTION ---
    logger.info("Setting up 'users' collection...")
    user_validator = {
        "$jsonSchema": {
            "bsonType": "object",
            "required": ["user_id", "phone", "language_preference", "is_active"],
            "properties": {
                "user_id": {"bsonType": "string"},
                "phone": {"bsonType": "string", "pattern": "^03\\d{2}-\\d{7}$"},
                "language_preference": {"enum": ["ur", "roman_ur", "en"]},
                "family_circle": {
                    "bsonType": "array",
                    "items": {
                        "bsonType": "object",
                        "properties": {
                            "name": {"bsonType": "string"},
                            "phone": {"bsonType": "string"},
                            "language": {"bsonType": "string"},
                            "relationship": {"bsonType": "string"}
                        }
                    }
                },
                "wallet_accounts": {
                    "bsonType": "array",
                    "items": {
                        "bsonType": "object",
                        "properties": {
                            "provider": {"enum": ["jazzcash", "easypaisa", "raast", "nayapay", "meezan"]},
                            "account_id": {"bsonType": "string"}
                        }
                    }
                },
                "is_active": {"bsonType": "bool"}
            }
        }
    }
    
    try:
        await db.create_collection("users", validator=user_validator)
    except Exception: pass
    await db.users.create_indexes([
        IndexModel([("phone", ASCENDING)], unique=True),
        IndexModel([("user_id", ASCENDING)], unique=True)
    ])

    # --- 2. SIGNALS COLLECTION ---
    logger.info("Setting up 'signals' collection...")
    try:
        await db.create_collection("signals")
    except Exception: pass
    await db.signals.create_indexes([
        IndexModel([("user_id", ASCENDING)]),
        IndexModel([("submitted_at", DESCENDING)])
    ])

    # --- 3. VERDICTS COLLECTION ---
    logger.info("Setting up 'verdicts' collection...")
    try:
        await db.create_collection("verdicts")
    except Exception: pass
    await db.verdicts.create_indexes([
        IndexModel([("signal_id", ASCENDING)]),
        IndexModel([("user_id", ASCENDING)]),
        IndexModel([("risk_level", ASCENDING)])
    ])

    # --- 4. FRAUD LEDGER COLLECTION ---
    logger.info("Setting up 'fraud_ledger' collection...")
    try:
        await db.create_collection("fraud_ledger")
    except Exception: pass
    await db.fraud_ledger.create_indexes([
        IndexModel([("sender_hash", ASCENDING)], unique=True),
        IndexModel([("url_hash", ASCENDING)]),
        IndexModel([("scam_pattern", TEXT)]),
        IndexModel([("hit_count", DESCENDING)])
    ])

    # --- 5. COMPLAINTS COLLECTION ---
    logger.info("Setting up 'complaints' collection...")
    try:
        await db.create_collection("complaints")
    except Exception: pass
    await db.complaints.create_indexes([
        IndexModel([("user_id", ASCENDING)]),
        IndexModel([("fia_ref", ASCENDING)], unique=True),
        IndexModel([("status", ASCENDING)])
    ])

    # --- 6. WALLET HOLDS COLLECTION ---
    logger.info("Setting up 'wallet_holds' collection...")
    try:
        await db.create_collection("wallet_holds")
    except Exception: pass
    await db.wallet_holds.create_indexes([
        IndexModel([("user_id", ASCENDING)]),
        IndexModel([("is_active", ASCENDING)]),
        IndexModel([("expires_at", ASCENDING)])
    ])

    # --- 7. TRACES COLLECTION ---
    logger.info("Setting up 'traces' collection...")
    try:
        await db.create_collection("traces")
    except Exception: pass
    await db.traces.create_indexes([
        IndexModel([("workplan_id", ASCENDING)], unique=True),
        IndexModel([("signal_id", ASCENDING)]),
        IndexModel([("created_at", DESCENDING)])
    ])

    # --- 8. ALERTS COLLECTION ---
    logger.info("Setting up 'alerts' collection...")
    try:
        await db.create_collection("alerts")
    except Exception: pass
    await db.alerts.create_indexes([
        IndexModel([("signal_id", ASCENDING)]),
        IndexModel([("sent_by_user_id", ASCENDING)])
    ])

    # --- SEED DATA ---
    logger.info("Inserting seed data for demo...")
    
    # Clean old seed data for a fresh setup
    await db.users.delete_many({})
    await db.fraud_ledger.delete_many({})

    # Seed Users
    users_seed = [
        {
            "user_id": str(uuid.uuid4()),
            "phone": "0312-3456789",
            "language_preference": "roman_ur",
            "family_circle": [{"name": "Beta Kamran", "phone": "0333-1234567", "language": "ur", "relationship": "son"}],
            "wallet_accounts": [{"provider": "meezan", "account_id": "PK77MEZN0001"}],
            "is_active": True,
            "created_at": datetime.now()
        },
        {
            "user_id": str(uuid.uuid4()),
            "phone": "0321-9876543",
            "language_preference": "ur",
            "family_circle": [],
            "wallet_accounts": [{"provider": "jazzcash", "account_id": "03219876543"}],
            "is_active": True,
            "created_at": datetime.now()
        }
    ]
    await db.users.insert_many(users_seed)

    # Seed Fraud Ledger (Pakistani Scam Patterns)
    ledger_seed = [
        {"sender_hash": "hash_jazz_help", "scam_pattern": "JazzCash Fake Helpline Calling from 0300 numbers", "scam_type": "fake_helpline", "hit_count": 47, "reported_by_count": 12},
        {"sender_hash": "hash_ep_otp", "scam_pattern": "EasyPaisa OTP request for 'Wrong Transfer' refund", "scam_type": "smishing", "hit_count": 23, "reported_by_count": 9},
        {"sender_hash": "hash_olx", "scam_pattern": "OLX Advance Payment for Bike or iPhone", "scam_type": "olx_fraud", "hit_count": 31, "reported_by_count": 15},
        {"sender_hash": "hash_wa_imp", "scam_pattern": "WhatsApp Impersonation 'Assalam o Alaikum, Main mushkil mein hoon'", "scam_type": "whatsapp_impersonation", "hit_count": 15, "reported_by_count": 6},
        {"sender_hash": "hash_loan", "scam_pattern": "Urgent Cash Loan App Blackmail/Extortion", "scam_type": "loan_extortion", "hit_count": 8, "reported_by_count": 4},
        {"sender_hash": "hash_benazir", "scam_pattern": "BISP / Benazir Income Support Program fake SMS", "scam_type": "smishing", "hit_count": 55, "reported_by_count": 20},
        {"sender_hash": "hash_jeeto", "scam_pattern": "Jeeto Pakistan / Inaami Scheme fraud message", "scam_type": "investment_scam", "hit_count": 12, "reported_by_count": 3},
        {"sender_hash": "hash_hbl", "scam_pattern": "HBL Account Blocked fake link phishing", "scam_type": "smishing", "hit_count": 19, "reported_by_count": 7},
        {"sender_hash": "hash_lottery", "scam_pattern": "Saudi Arabia Lottery fake WhatsApp call", "scam_type": "investment_scam", "hit_count": 5, "reported_by_count": 2},
        {"sender_hash": "hash_fbr", "scam_pattern": "FBR Tax Refund fake website scam", "scam_type": "unknown", "hit_count": 9, "reported_by_count": 4}
    ]
    await db.fraud_ledger.insert_many(ledger_seed)

    await verify_setup(db)
    logger.info("NadraGuard DB setup complete.")

async def verify_setup(db):
    logger.info("--- VERIFICATION SUMMARY ---")
    collections = await db.list_collection_names()
    for col in collections:
        count = await db[col].count_documents({})
        indexes = await db[col].list_indexes().to_list(length=100)
        logger.info(f"Collection: {col.ljust(15)} | Documents: {str(count).ljust(5)} | Indexes: {len(indexes)}")

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