import asyncio
from motor.motor_asyncio import AsyncIOMotorClient
import os
from app.core.config import Config


MONGODB_URI = Config.MONGODB_URI
DB_NAME = Config.MONGODB_DATABASE

async def main():
    client = AsyncIOMotorClient(MONGODB_URI, tz_aware=True)
    database = client[DB_NAME]

    print("[ASYNC_DATABASE] Starting collection and index creation...")
    
    # List existing collections before creation
    existing_collections = await database.list_collection_names()
    print(f"[ASYNC_DATABASE] Existing collections before creation: {existing_collections}")
    
    # Models collection
    print("[ASYNC_DATABASE] Creating models collection...")
    models_collection = database["models"]
    await models_collection.create_index([("model_id", 1)], unique=True)
    await models_collection.create_index([("user_id", 1), ("model_name", 1)])
    await models_collection.create_index([("user_id", 1), ("tag", 1)])
    await models_collection.create_index([("status", 1)])
    await models_collection.create_index([("organization_id", 1)])
    await models_collection.create_index([("department_id", 1)])
    await models_collection.create_index([("created_at", -1)])
    print("[ASYNC_DATABASE] Models collection and indexes created")
    
    # Training runs collection
    print("[ASYNC_DATABASE] Creating training_runs collection...")
    training_runs_collection = database["training_runs"]
    await training_runs_collection.create_index([("training_id", 1)], unique=True)
    await training_runs_collection.create_index([("model_id", 1)])
    await training_runs_collection.create_index([("user_id", 1)])
    await training_runs_collection.create_index([("status", 1)])
    await training_runs_collection.create_index([("created_at", -1)])
    await training_runs_collection.create_index([("compliance_type", 1)])
    print("[ASYNC_DATABASE] Training runs collection and indexes created")
    
    # Validation runs collection
    print("[ASYNC_DATABASE] Creating validation_runs collection...")
    validation_runs_collection = database["validation_runs"]
    await validation_runs_collection.create_index([("validation_id", 1)], unique=True)
    await validation_runs_collection.create_index([("training_id", 1)])
    await validation_runs_collection.create_index([("model_id", 1)])
    await validation_runs_collection.create_index([("model_name", 1)])
    await validation_runs_collection.create_index([("user_id", 1)])
    await validation_runs_collection.create_index([("created_at", -1)])
    await validation_runs_collection.create_index([("compliance_type", 1)])
    print("[ASYNC_DATABASE] Validation runs collection and indexes created")
    
    # Prediction runs collection
    print("[ASYNC_DATABASE] Creating prediction_runs collection...")
    prediction_runs_collection = database["prediction_runs"]
    await prediction_runs_collection.create_index([("prediction_id", 1)], unique=True)
    await prediction_runs_collection.create_index([("training_id", 1)])
    await prediction_runs_collection.create_index([("model_id", 1)])
    await prediction_runs_collection.create_index([("model_name", 1)])
    await prediction_runs_collection.create_index([("user_id", 1)])
    await prediction_runs_collection.create_index([("created_at", -1)])
    await prediction_runs_collection.create_index([("compliance_type", 1)])
    print("[ASYNC_DATABASE] Prediction runs collection and indexes created")
    
    # System status collection
    print("[ASYNC_DATABASE] Creating system_status collection...")
    status_collection = database["system_status"]
    await status_collection.create_index([("training_id", 1)], unique=True)
    await status_collection.create_index([("user_id", 1)])
    await status_collection.create_index([("status", 1)])
    await status_collection.create_index([("created_at", -1)])
    print("[ASYNC_DATABASE] System status collection and indexes created")
    
    # -----------------------------------------------------------------
    # Organizations collection
    # -----------------------------------------------------------------
    print("[ASYNC_DATABASE] Creating organizations collection...")
    organizations = database["organizations"]
    await organizations.create_index([("name", 1)])
    print("[ASYNC_DATABASE] Organizations collection and indexes created")

    # -----------------------------------------------------------------
    # Departments collection
    # -----------------------------------------------------------------
    print("[ASYNC_DATABASE] Creating departments collection...")
    departments = database["departments"]
    await departments.create_index([("name", 1)])
    # Foreign key style index to organization
    await departments.create_index([("organization_id", 1)])
    print("[ASYNC_DATABASE] Departments collection and indexes created")

    # -----------------------------------------------------------------
    # Users collection
    # -----------------------------------------------------------------
    print("[ASYNC_DATABASE] Creating users collection...")
    users = database["users"]
    await users.create_index([("username", 1)])
    # Foreign key style references
    await users.create_index([("organization_id", 1)])
    await users.create_index([("department_id", 1)])
    # Admin type filter
    await users.create_index([("admin_type", 1)])
    print("[ASYNC_DATABASE] Users collection and indexes created")

    # -----------------------------------------------------------------
    # Metadata collection
    # -----------------------------------------------------------------
    print("[ASYNC_DATABASE] Creating Metadata collection...")
    meta_data = database["metadata"]
    # await meta_data.create_index([("metadata_id", 1)])
    # await meta_data.create_index([("feature_data", 1)])
    print("[ASYNC_DATABASE] Metadata collection and indexes created")
    
    # Verify collections were created
    final_collections = await database.list_collection_names()
    print(f"[ASYNC_DATABASE] Final collections after creation: {final_collections}")
    
    print("[ASYNC_DATABASE] All collections and indexes created successfully")

    print("[INDEX] All indexes created successfully.")
    client.close()


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