"""
Migration script: Assign organization_id to all existing documents.

Run once after deploying the multi-tenant backend changes.
Usage:
    python migrate_to_multitenant.py
"""
import asyncio
import os
import uuid
from datetime import datetime

from motor.motor_asyncio import AsyncIOMotorClient
from dotenv import load_dotenv

load_dotenv()


COLLECTIONS_TO_MIGRATE = [
    "devices",
    "scans",
    "alerts",
    "device_metrics",
    "backups",
    "assets",
    "asset_movements",
    "firmwares",
    "netflow_records",
    "gis_cables",
    "gis_junctions",
    "gis_fibers",
    "gis_incidents",
    "customers",
    "invoices",
    "payments",
    "packages",
    "chat_history",
    "knowledge_base",
]


async def migrate():
    mongo_url = os.environ["MONGO_URL"]
    db_name = os.environ["DB_NAME"]
    client = AsyncIOMotorClient(mongo_url)
    db = client[db_name]

    org = await db.organizations.find_one({}, sort=[("created_at", 1)])

    if not org:
        org_id = str(uuid.uuid4())
        org_doc = {
            "id": org_id,
            "name": "Default Organization",
            "plan": "free",
            "created_at": datetime.utcnow(),
            "features_enabled": [],
            "max_users": 50,
            "max_devices": 500,
        }
        await db.organizations.insert_one(org_doc)
        print(f"Created default organization: {org_id}")
    else:
        org_id = org["id"]
        print(f"Using existing organization: {org_id}")

    users_result = await db.users.update_many(
        {"organization_id": {"$exists": False}},
        {"$set": {"organization_id": org_id}},
    )
    print(f"Users updated: {users_result.modified_count}")

    first_user = await db.users.find_one(
        {"organization_id": org_id},
        sort=[("created_at", 1)],
    )
    if first_user and not first_user.get("is_org_owner"):
        await db.users.update_one(
            {"id": first_user["id"]},
            {"$set": {"is_org_owner": True}},
        )
        print(f"Marked {first_user.get('email', first_user['id'])} as org owner")

    for collection_name in COLLECTIONS_TO_MIGRATE:
        collection = db[collection_name]
        result = await collection.update_many(
            {"organization_id": {"$exists": False}},
            {"$set": {"organization_id": org_id}},
        )
        print(f"  {collection_name}: {result.modified_count} docs updated")

    print("\nMigration complete.")
    client.close()


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