# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""
Migration: Root-Mandant bereinigen.
Moves all end-user data from Root mandate shared instances to own mandates.
Called once from bootstrap, sets a DB flag to prevent re-execution.
"""

import logging
from typing import Optional, List, Dict, Any

logger = logging.getLogger(__name__)

_MIGRATION_FLAG_KEY = "migration_root_users_completed"

_DATA_TABLES = [
    "ChatWorkflow",
    "FileItem",
    "DataSource",
    "DataNeutralizerAttributes",
    "FileContentIndex",
]


def _isMigrationCompleted(db) -> bool:
    """Check if migration has already been executed."""
    try:
        from modules.datamodels.datamodelUam import Mandate
        records = db.getRecordset(Mandate, recordFilter={"name": _MIGRATION_FLAG_KEY})
        return len(records) > 0
    except Exception:
        return False


def _setMigrationCompleted(db) -> None:
    """Set flag that migration is completed (uses a settings-like record)."""
    if _isMigrationCompleted(db):
        return
    try:
        from modules.datamodels.datamodelUam import Mandate
        flag = Mandate(name=_MIGRATION_FLAG_KEY, label="Migration completed", enabled=False, isSystem=True)
        db.recordCreate(Mandate, flag)
        logger.info("Migration flag set: root user migration completed")
    except Exception as e:
        logger.error(f"Failed to set migration flag: {e}")


def _findOrCreateTargetInstance(db, featureInterface, featureCode: str, targetMandateId: str, rootInstance: dict) -> dict:
    """Find existing or create new FeatureInstance in target mandate. Idempotent."""
    from modules.datamodels.datamodelFeatures import FeatureInstance

    existing = db.getRecordset(FeatureInstance, recordFilter={
        "featureCode": featureCode,
        "mandateId": targetMandateId,
    })
    if existing:
        logger.debug(f"Target instance already exists for {featureCode} in mandate {targetMandateId}")
        return existing[0]

    label = rootInstance.get("label") or featureCode
    instance = featureInterface.createFeatureInstance(
        featureCode=featureCode,
        mandateId=targetMandateId,
        label=label,
        enabled=True,
        copyTemplateRoles=True,
    )
    if isinstance(instance, dict):
        return instance
    return instance.model_dump() if hasattr(instance, "model_dump") else {"id": instance.id}


def _migrateDataRecords(db, oldInstanceId: str, newInstanceId: str, userId: str) -> int:
    """Bulk-update featureInstanceId on all data tables for records owned by userId."""
    totalMigrated = 0
    db._ensure_connection()
    for tableName in _DATA_TABLES:
        try:
            with db.connection.cursor() as cursor:
                cursor.execute(
                    f'UPDATE "{tableName}" '
                    f'SET "featureInstanceId" = %s '
                    f'WHERE "featureInstanceId" = %s AND "sysCreatedBy" = %s',
                    (newInstanceId, oldInstanceId, userId),
                )
                count = cursor.rowcount
            db.connection.commit()
            if count > 0:
                logger.info(f"  Migrated {count} rows in {tableName}: {oldInstanceId} -> {newInstanceId}")
                totalMigrated += count
        except Exception as e:
            try:
                db.connection.rollback()
            except Exception:
                pass
            logger.debug(f"  Table {tableName} skipped (may not exist or no matching column): {e}")
    return totalMigrated


def _grantFeatureAccess(db, userId: str, featureInstanceId: str) -> dict:
    """Create FeatureAccess + admin role on a feature instance. Idempotent."""
    from modules.datamodels.datamodelMembership import FeatureAccess, FeatureAccessRole
    from modules.datamodels.datamodelRbac import Role

    existing = db.getRecordset(FeatureAccess, recordFilter={
        "userId": userId,
        "featureInstanceId": featureInstanceId,
    })
    if existing:
        logger.debug(f"FeatureAccess already exists for user {userId} on instance {featureInstanceId}")
        return existing[0]

    fa = FeatureAccess(userId=userId, featureInstanceId=featureInstanceId, enabled=True)
    createdFa = db.recordCreate(FeatureAccess, fa.model_dump())
    if not createdFa:
        logger.warning(f"Failed to create FeatureAccess for user {userId} on instance {featureInstanceId}")
        return {}

    instanceRoles = db.getRecordset(Role, recordFilter={"featureInstanceId": featureInstanceId})
    adminRoleId = None
    for r in instanceRoles:
        roleLabel = (r.get("roleLabel") or "").lower()
        if roleLabel.endswith("-admin"):
            adminRoleId = r.get("id")
            break
    if not adminRoleId:
        raise ValueError(
            f"No feature-specific admin role for instance {featureInstanceId}. "
            f"Cannot create FeatureAccess without role — even in migration context."
        )
    far = FeatureAccessRole(featureAccessId=createdFa["id"], roleId=adminRoleId)
    db.recordCreate(FeatureAccessRole, far.model_dump())

    return createdFa


def migrateRootUsers(db, dryRun: bool = False) -> dict:
    """
    Migrate all end-user feature data from Root mandate to personal mandates.
    
    Algorithm:
    STEP 1: For each user with FeatureAccess on Root instances:
      - If user has own mandate: target = existing mandate
      - If not: create personal mandate via _provisionMandateForUser
      - For each FeatureAccess: create new instance in target, migrate data, transfer access
    
    STEP 2: Clean up Root:
      - Delete all FeatureInstances in Root
      - Remove UserMandate for non-sysadmin users
    
    Args:
        db: Database connector
        dryRun: If True, log actions without making changes
        
    Returns:
        Summary dict with migration statistics
    """
    if _isMigrationCompleted(db):
        logger.info("Root user migration already completed, skipping")
        return {"status": "already_completed"}

    from modules.datamodels.datamodelUam import Mandate, User, UserInDB
    from modules.datamodels.datamodelMembership import (
        UserMandate, UserMandateRole, FeatureAccess, FeatureAccessRole,
    )
    from modules.datamodels.datamodelFeatures import FeatureInstance
    from modules.interfaces.interfaceDbApp import getRootInterface
    from modules.interfaces.interfaceFeatures import getFeatureInterface

    rootInterface = getRootInterface()
    featureInterface = getFeatureInterface(db)
    stats = {
        "usersProcessed": 0,
        "mandatesCreated": 0,
        "instancesMigrated": 0,
        "dataRowsMigrated": 0,
        "rootInstancesDeleted": 0,
        "rootMembershipsRemoved": 0,
        "dryRun": dryRun,
    }

    # Find root mandate
    rootMandates = db.getRecordset(Mandate, recordFilter={"name": "root", "isSystem": True})
    if not rootMandates:
        logger.warning("No root mandate found, nothing to migrate")
        return {"status": "no_root_mandate"}
    rootMandateId = rootMandates[0].get("id")

    # Get all feature instances in root
    rootInstances = db.getRecordset(FeatureInstance, recordFilter={"mandateId": rootMandateId})
    if not rootInstances:
        logger.info("No feature instances in root mandate, nothing to migrate")
        if not dryRun:
            _setMigrationCompleted(db)
        return {"status": "no_instances", **stats}

    # Get all FeatureAccess on root instances
    rootInstanceIds = {inst.get("id") for inst in rootInstances}

    # Collect unique users with access on root instances
    usersToMigrate = {}
    for instanceId in rootInstanceIds:
        accesses = db.getRecordset(FeatureAccess, recordFilter={"featureInstanceId": instanceId})
        for access in accesses:
            userId = access.get("userId")
            if userId not in usersToMigrate:
                usersToMigrate[userId] = []
            usersToMigrate[userId].append({
                "featureAccessId": access.get("id"),
                "featureInstanceId": instanceId,
            })

    logger.info(f"Migration: {len(usersToMigrate)} users with {sum(len(v) for v in usersToMigrate.values())} accesses on {len(rootInstances)} root instances")

    # STEP 1: Migrate users
    for userId, accessList in usersToMigrate.items():
        try:
            # Find user
            users = db.getRecordset(UserInDB, recordFilter={"id": userId})
            if not users:
                logger.warning(f"User {userId} not found, skipping")
                continue
            user = users[0]
            username = user.get("username", "unknown")

            # Check if user has own non-root mandate
            userMandates = db.getRecordset(UserMandate, recordFilter={"userId": userId, "enabled": True})
            targetMandateId = None
            for um in userMandates:
                mid = um.get("mandateId")
                if mid != rootMandateId:
                    targetMandateId = mid
                    break

            if not targetMandateId:
                # Create personal mandate
                if dryRun:
                    logger.info(f"[DRY RUN] Would create personal mandate for user {username}")
                    stats["mandatesCreated"] += 1
                else:
                    try:
                        result = rootInterface._provisionMandateForUser(
                            userId=userId,
                            mandateName=f"Home {username}",
                            planKey="TRIAL_14D",
                        )
                        targetMandateId = result["mandateId"]
                        stats["mandatesCreated"] += 1
                        logger.info(f"Created personal mandate {targetMandateId} for user {username}")
                    except Exception as e:
                        logger.error(f"Failed to create mandate for user {username}: {e}")
                        continue

            # Migrate each FeatureAccess
            for accessInfo in accessList:
                oldInstanceId = accessInfo["featureInstanceId"]
                oldAccessId = accessInfo["featureAccessId"]

                # Find the root instance details
                instRecords = db.getRecordset(FeatureInstance, recordFilter={"id": oldInstanceId})
                if not instRecords:
                    continue
                featureCode = instRecords[0].get("featureCode")

                if dryRun:
                    logger.info(f"[DRY RUN] Would migrate {featureCode} for {username} to mandate {targetMandateId}")
                    stats["instancesMigrated"] += 1
                else:
                    targetInstance = _findOrCreateTargetInstance(
                        db, featureInterface, featureCode, targetMandateId, instRecords[0],
                    )
                    newInstanceId = targetInstance.get("id")
                    if not newInstanceId:
                        logger.error(f"Failed to obtain target instance for {featureCode} in mandate {targetMandateId}")
                        continue

                    migratedCount = _migrateDataRecords(db, oldInstanceId, newInstanceId, userId)

                    _grantFeatureAccess(db, userId, newInstanceId)

                    try:
                        db.recordDelete(FeatureAccess, oldAccessId)
                    except Exception as delErr:
                        logger.warning(f"Could not remove old FeatureAccess {oldAccessId}: {delErr}")

                    logger.info(
                        f"Migrated {featureCode} for {username}: "
                        f"instance {oldInstanceId} -> {newInstanceId}, {migratedCount} data rows moved"
                    )
                    stats["instancesMigrated"] += 1
                    stats["dataRowsMigrated"] += migratedCount

            stats["usersProcessed"] += 1

        except Exception as e:
            logger.error(f"Error migrating user {userId}: {e}")

    # STEP 2: Clean up root
    if not dryRun:
        # Delete all feature instances in root
        for inst in rootInstances:
            instId = inst.get("id")
            try:
                # First delete all FeatureAccess on this instance
                accesses = db.getRecordset(FeatureAccess, recordFilter={"featureInstanceId": instId})
                for access in accesses:
                    db.recordDelete(FeatureAccess, access.get("id"))
                db.recordDelete(FeatureInstance, instId)
                stats["rootInstancesDeleted"] += 1
            except Exception as e:
                logger.error(f"Error deleting root instance {instId}: {e}")

        # Remove non-sysadmin users from root mandate
        rootMembers = db.getRecordset(UserMandate, recordFilter={"mandateId": rootMandateId})
        for membership in rootMembers:
            membUserId = membership.get("userId")
            userRecords = db.getRecordset(UserInDB, recordFilter={"id": membUserId})
            if userRecords and userRecords[0].get("isSysAdmin"):
                continue
            try:
                db.recordDelete(UserMandate, membership.get("id"))
                stats["rootMembershipsRemoved"] += 1
            except Exception as e:
                logger.error(f"Error removing root membership for {membUserId}: {e}")

        _setMigrationCompleted(db)

    logger.info(f"Migration completed: {stats}")
    return {"status": "completed", **stats}