# -*- coding: utf-8 -*-
"""Location: ./mcpgateway/bootstrap_db.py
Copyright 2025
SPDX-License-Identifier: Apache-2.0
Authors: Madhav Kandukuri

Database bootstrap/upgrade entry-point for MCP Gateway.
The script:

1. Creates a synchronous SQLAlchemy ``Engine`` from ``settings.database_url``.
2. Looks for an *alembic.ini* two levels up from this file to drive migrations.
3. Applies Alembic migrations (``alembic upgrade head``) to create or update the schema.
4. Runs post-upgrade normalization tasks and bootstraps admin/roles as configured.
5. Logs a **"Database ready"** message on success.

It is intended to be invoked via ``python3 -m mcpgateway.bootstrap_db`` or
directly with ``python3 mcpgateway/bootstrap_db.py``.

Examples:
    >>> from mcpgateway.bootstrap_db import logging_service, logger
    >>> logging_service is not None
    True
    >>> logger is not None
    True
    >>> hasattr(logger, 'info')
    True
    >>> from mcpgateway.bootstrap_db import Base
    >>> hasattr(Base, 'metadata')
    True
"""

# Standard
import asyncio
from contextlib import contextmanager
from importlib.resources import files
import json
import os
from pathlib import Path
import tempfile
from typing import cast

# Third-Party
from alembic import command
from alembic.config import Config
from filelock import FileLock
from sqlalchemy import create_engine, inspect, text
from sqlalchemy.engine import Connection
from sqlalchemy.orm import Session

# First-Party
from mcpgateway.config import settings
from mcpgateway.db import A2AAgent, Base, EmailTeam, EmailUser, Gateway, Prompt, Resource, Server, Tool
from mcpgateway.services.logging_service import LoggingService

# Migration lock to prevent concurrent migrations from multiple workers
_MIGRATION_LOCK_PATH = os.path.join(tempfile.gettempdir(), "mcpgateway_migration.lock")
_MIGRATION_LOCK_TIMEOUT = 300  # seconds to wait for lock (5 minutes for slow migrations)

# Initialize logging service first
logging_service = LoggingService()
logger = logging_service.get_logger(__name__)


def _column_exists(inspector, table_name: str, column_name: str) -> bool:
    """Check whether a table has a specific column.

    Args:
        inspector: SQLAlchemy inspector for the active connection.
        table_name: Table name to inspect.
        column_name: Column name to check.

    Returns:
        True if the column exists, otherwise False.
    """
    try:
        return any(col["name"] == column_name for col in inspector.get_columns(table_name))
    except Exception:
        return False


def _schema_looks_current(inspector) -> bool:
    """Best-effort check for unversioned databases that already match current schema.

    Args:
        inspector: SQLAlchemy inspector for the active connection.

    Returns:
        True when expected columns exist for a recent schema version.
    """
    return (
        _column_exists(inspector, "tools", "display_name")
        and _column_exists(inspector, "gateways", "oauth_config")
        and _column_exists(inspector, "prompts", "custom_name")
        and _column_exists(inspector, "sso_providers", "jwks_uri")
    )


@contextmanager
def advisory_lock(conn: Connection):
    """
    Acquire a distributed advisory lock to serialize migrations across multiple instances.

    Behavior depends on the database backend:
    - Postgres: Uses `pg_advisory_lock` (blocking)
    - MySQL: Uses `GET_LOCK` (blocking with timeout)
    - SQLite: Fallback to local `FileLock`

    Args:
        conn: Active SQLAlchemy connection

    Yields:
        None

    Raises:
        TimeoutError: If the lock cannot be acquired within the timeout period
    """
    dialect = conn.dialect.name
    lock_id = "mcpgateway_migration"
    # Postgres requires a BIGINT lock ID (arbitrary hash of the string)
    pg_lock_id = 42424242424242

    if dialect == "postgresql":
        logger.info("Acquiring Postgres advisory lock...")
        conn.execute(text(f"SELECT pg_advisory_lock({pg_lock_id})"))
        try:
            yield
        finally:
            logger.info("Releasing Postgres advisory lock...")
            conn.execute(text(f"SELECT pg_advisory_unlock({pg_lock_id})"))

    elif dialect in ["mysql", "mariadb"]:
        logger.info("Acquiring MySQL advisory lock...")
        # GET_LOCK returns 1 if successful, 0 if timed out, NULL on error
        result = conn.execute(text(f"SELECT GET_LOCK('{lock_id}', {_MIGRATION_LOCK_TIMEOUT})")).scalar()
        if result != 1:
            raise TimeoutError(f"Could not acquire MySQL lock '{lock_id}' within {_MIGRATION_LOCK_TIMEOUT}s")
        try:
            yield
        finally:
            logger.info("Releasing MySQL advisory lock...")
            conn.execute(text(f"SELECT RELEASE_LOCK('{lock_id}')"))

    else:
        # Fallback for SQLite (single-host/container) or other DBs
        logger.info(f"Using FileLock fallback for {dialect}...")
        file_lock = FileLock(_MIGRATION_LOCK_PATH, timeout=_MIGRATION_LOCK_TIMEOUT)
        with file_lock:
            yield


async def bootstrap_admin_user(conn: Connection) -> None:
    """
    Bootstrap the platform admin user from environment variables.

    Creates the admin user if email authentication is enabled and the user doesn't exist.
    Also creates a personal team for the admin user if auto-creation is enabled.

    Args:
        conn: Active SQLAlchemy connection
    """
    if not settings.email_auth_enabled:
        logger.info("Email authentication disabled - skipping admin user bootstrap")
        return

    try:
        # Import services here to avoid circular imports
        # First-Party
        from mcpgateway.services.email_auth_service import EmailAuthService  # pylint: disable=import-outside-toplevel

        # Use session bound to the locked connection
        with Session(bind=conn) as db:
            auth_service = EmailAuthService(db)

            # Check if admin user already exists
            existing_user = await auth_service.get_user_by_email(settings.platform_admin_email)
            if existing_user:
                logger.info(f"Admin user {settings.platform_admin_email} already exists - skipping creation")
                return

            # Create admin user
            logger.info(f"Creating platform admin user: {settings.platform_admin_email}")
            admin_user = await auth_service.create_platform_admin(
                email=settings.platform_admin_email,
                password=settings.platform_admin_password.get_secret_value(),
                full_name=settings.platform_admin_full_name,
            )

            # Mark admin user as email verified and require password change on first login
            # First-Party
            from mcpgateway.db import utc_now  # pylint: disable=import-outside-toplevel

            admin_user.email_verified_at = utc_now()
            # Respect configuration: only require password change on bootstrap when enabled
            if getattr(settings, "password_change_enforcement_enabled", True) and getattr(settings, "admin_require_password_change_on_bootstrap", True):
                admin_user.password_change_required = True  # Force admin to change default password
            try:
                admin_user.password_changed_at = utc_now()
            except Exception as exc:
                logger.debug("Failed to set admin password_changed_at: %s", exc)
            db.commit()

            # Personal team is automatically created during user creation if enabled
            if settings.auto_create_personal_teams:
                logger.info("Personal team automatically created for admin user")

            db.commit()
            logger.info(f"Platform admin user created successfully: {settings.platform_admin_email}")

    except Exception as e:
        logger.error(f"Failed to bootstrap admin user: {e}")
        # Don't fail the entire bootstrap process if admin user creation fails
        return


async def bootstrap_default_roles(conn: Connection) -> None:
    """Bootstrap default system roles and assign them to admin user.

    Creates essential RBAC roles and assigns administrative privileges
    to the platform admin user.

    Args:
        conn: Active SQLAlchemy connection
    """
    if not settings.email_auth_enabled:
        logger.info("Email authentication disabled - skipping default roles bootstrap")
        return

    try:
        # First-Party
        from mcpgateway.services.email_auth_service import EmailAuthService  # pylint: disable=import-outside-toplevel
        from mcpgateway.services.role_service import RoleService  # pylint: disable=import-outside-toplevel

        # Use session bound to the locked connection
        with Session(bind=conn) as db:
            role_service = RoleService(db)
            auth_service = EmailAuthService(db)

            # Check if admin user exists
            admin_user = await auth_service.get_user_by_email(settings.platform_admin_email)
            if not admin_user:
                logger.info("Admin user not found - skipping role assignment")
                return

            # Default system roles to create
            default_roles = [
                {"name": "platform_admin", "description": "Platform administrator with all permissions", "scope": "global", "permissions": ["*"], "is_system_role": True},  # All permissions
                {
                    "name": "team_admin",
                    "description": "Team administrator with team management permissions",
                    "scope": "team",
                    "permissions": [
                        "admin.dashboard",
                        "gateways.read",
                        "servers.read",
                        "servers.use",
                        "teams.read",
                        "teams.update",
                        "teams.join",
                        "teams.delete",
                        "teams.manage_members",
                        "tools.read",
                        "tools.execute",
                        "resources.read",
                        "prompts.read",
                        "a2a.read",
                        "gateways.create",
                        "servers.create",
                        "tools.create",
                        "resources.create",
                        "prompts.create",
                        "a2a.create",
                        "gateways.update",
                        "servers.update",
                        "tools.update",
                        "resources.update",
                        "prompts.update",
                        "a2a.update",
                        "gateways.delete",
                        "servers.delete",
                        "tools.delete",
                        "resources.delete",
                        "prompts.delete",
                        "a2a.delete",
                        "a2a.invoke",
                        "tokens.create",
                        "tokens.read",
                        "tokens.update",
                        "tokens.revoke",
                    ],
                    "is_system_role": True,
                },
                {
                    "name": "developer",
                    "description": "Developer with tool and resource access",
                    "scope": "team",
                    "permissions": [
                        "admin.dashboard",
                        "gateways.read",
                        "servers.read",
                        "servers.use",
                        "teams.join",
                        "tools.read",
                        "tools.execute",
                        "resources.read",
                        "prompts.read",
                        "a2a.read",
                        "gateways.create",
                        "servers.create",
                        "tools.create",
                        "resources.create",
                        "prompts.create",
                        "a2a.create",
                        "gateways.update",
                        "servers.update",
                        "tools.update",
                        "resources.update",
                        "prompts.update",
                        "a2a.update",
                        "gateways.delete",
                        "servers.delete",
                        "tools.delete",
                        "resources.delete",
                        "prompts.delete",
                        "a2a.delete",
                        "a2a.invoke",
                    ],
                    "is_system_role": True,
                },
                {
                    "name": "viewer",
                    "description": "Read-only access to resources and admin UI",
                    "scope": "team",
                    "permissions": ["admin.dashboard", "gateways.read", "servers.read", "teams.join", "tools.read", "resources.read", "prompts.read", "a2a.read"],
                    "is_system_role": True,
                },
                {
                    "name": "platform_viewer",
                    "description": "Read-only access to resources and admin UI",
                    "scope": "global",
                    "permissions": ["admin.dashboard", "gateways.read", "servers.read", "teams.join", "tools.read", "resources.read", "prompts.read", "a2a.read"],
                    "is_system_role": True,
                },
            ]

            # Logic to add additional default roles from a json file
            if settings.mcpgateway_bootstrap_roles_in_db_enabled:
                try:
                    additional_default_roles_path = Path(settings.mcpgateway_bootstrap_roles_in_db_file)
                    # Try multiple locations for the mcpgateway_bootstrap_roles_in_db_file file
                    if not additional_default_roles_path.is_absolute():
                        # Try current directory first
                        if not additional_default_roles_path.exists():
                            # Try project root (mcpgateway/bootstrap_db.py -> parent.parent = repo root)
                            additional_default_roles_path = Path(__file__).resolve().parent.parent / settings.mcpgateway_bootstrap_roles_in_db_file

                    if not additional_default_roles_path.exists():
                        logger.warning(f"Additional roles file not found. Searched: CWD/{settings.mcpgateway_bootstrap_roles_in_db_file}, {additional_default_roles_path}")
                    else:
                        with open(additional_default_roles_path, "r", encoding="utf-8") as f:
                            additional_default_roles_data = json.load(f)

                        # Validate JSON structure: must be a list of dicts with required keys
                        required_keys = {"name", "scope", "permissions"}
                        if not isinstance(additional_default_roles_data, list):
                            logger.error(f"Additional roles file must contain a JSON array, got {type(additional_default_roles_data).__name__}")
                        else:
                            valid_roles = []
                            for idx, role in enumerate(additional_default_roles_data):
                                if not isinstance(role, dict):
                                    logger.warning(f"Skipping invalid role at index {idx}: expected dict, got {type(role).__name__}")
                                    continue
                                missing_keys = required_keys - set(role.keys())
                                if missing_keys:
                                    role_name = role.get("name", f"<index {idx}>")
                                    logger.warning(f"Skipping role '{role_name}': missing required keys {missing_keys}")
                                    continue
                                valid_roles.append(role)

                            if valid_roles:
                                default_roles.extend(valid_roles)
                                logger.info(f"Added {len(valid_roles)} additional roles to default roles in bootstrap db")
                            elif additional_default_roles_data:
                                logger.warning("No valid roles found in additional roles file")
                except Exception as e:
                    logger.error(f"Failed to load mcpgateway_bootstrap_roles_in_db_file: {e}")

            # Create default roles
            created_roles = []
            for role_def in default_roles:
                try:
                    # Check if role already exists
                    existing_role = await role_service.get_role_by_name(str(role_def["name"]), str(role_def["scope"]))
                    if existing_role:
                        logger.info(f"System role {role_def['name']} already exists - skipping")
                        created_roles.append(existing_role)
                        continue

                    # Create the role (description and is_system_role are optional)
                    role = await role_service.create_role(
                        name=str(role_def["name"]),
                        description=str(role_def.get("description", "")),
                        scope=str(role_def["scope"]),
                        permissions=cast(list[str], role_def["permissions"]),
                        created_by=settings.platform_admin_email,
                        is_system_role=bool(role_def.get("is_system_role", False)),
                    )
                    created_roles.append(role)
                    logger.info(f"Created system role: {role.name}")

                except Exception as e:
                    logger.error(f"Failed to create role {role_def['name']}: {e}")
                    continue

            # Assign platform_admin role to admin user
            platform_admin_role = next((r for r in created_roles if r.name == "platform_admin"), None)
            if not platform_admin_role:
                # Role not in created_roles (creation may have failed) — look up from DB as fallback
                platform_admin_role = await role_service.get_role_by_name("platform_admin", "global")
            if platform_admin_role:
                try:
                    # Check if assignment already exists
                    existing_assignment = await role_service.get_user_role_assignment(user_email=admin_user.email, role_id=platform_admin_role.id, scope="global", scope_id=None)

                    if not existing_assignment or not existing_assignment.is_active:
                        await role_service.assign_role_to_user(user_email=admin_user.email, role_id=platform_admin_role.id, scope="global", scope_id=None, granted_by=admin_user.email)
                        logger.info(f"Assigned platform_admin role to {admin_user.email}")
                    else:
                        logger.info("Admin user already has platform_admin role")

                except Exception as e:
                    logger.error(f"Failed to assign platform_admin role to {admin_user.email}: {e}. Admin UI routes using allow_admin_bypass=False will return 403.")
            else:
                logger.error(f"platform_admin role not found — could not assign to {admin_user.email}. Admin UI routes using allow_admin_bypass=False will return 403.")

            logger.info("Default RBAC roles bootstrap completed successfully")

    except Exception as e:
        logger.error(f"Failed to bootstrap default roles: {e}")
        # Don't fail the entire bootstrap process if role creation fails
        return


def normalize_team_visibility(conn: Connection) -> int:
    """Normalize team visibility values to the supported set {private, public}.

    Any team with an unsupported visibility (e.g., 'team') is set to 'private'.

    Args:
        conn: Active SQLAlchemy connection

    Returns:
        int: Number of teams updated
    """
    try:
        # Use session bound to the locked connection
        with Session(bind=conn) as db:
            # Find teams with invalid visibility
            invalid = db.query(EmailTeam).filter(EmailTeam.visibility.notin_(["private", "public"]))
            count = 0
            for team in invalid.all():
                old = team.visibility
                team.visibility = "private"
                count += 1
                logger.info(f"Normalized team visibility: id={team.id} {old} -> private")
            if count:
                db.commit()
            return count
    except Exception as e:
        logger.error(f"Failed to normalize team visibility: {e}")
        return 0


async def bootstrap_resource_assignments(conn: Connection) -> None:
    """Assign orphaned resources to the platform admin's personal team.

    This ensures existing resources (from pre-multitenancy versions) are
    visible in the new team-based UI by assigning them to the admin's
    personal team with public visibility.

    Args:
        conn: Active SQLAlchemy connection
    """
    if not settings.email_auth_enabled:
        logger.info("Email authentication disabled - skipping resource assignment")
        return

    try:
        # Use session bound to the locked connection
        with Session(bind=conn) as db:
            # Find admin user and their personal team
            admin_user = db.query(EmailUser).filter(EmailUser.email == settings.platform_admin_email, EmailUser.is_admin.is_(True)).first()

            if not admin_user:
                logger.warning("Admin user not found - skipping resource assignment")
                return

            personal_team = admin_user.get_personal_team()
            if not personal_team:
                logger.warning("Admin personal team not found - skipping resource assignment")
                return

            logger.info(f"Assigning orphaned resources to admin team: {personal_team.name}")

            # Resource types to process
            resource_types = [("servers", Server), ("tools", Tool), ("resources", Resource), ("prompts", Prompt), ("gateways", Gateway), ("a2a_agents", A2AAgent)]

            total_assigned = 0

            for resource_name, resource_model in resource_types:
                try:
                    # Find unassigned resources
                    unassigned = db.query(resource_model).filter((resource_model.team_id.is_(None)) | (resource_model.owner_email.is_(None)) | (resource_model.visibility.is_(None))).all()

                    if unassigned:
                        logger.info(f"Assigning {len(unassigned)} orphaned {resource_name} to admin team")

                        for resource in unassigned:
                            resource.team_id = personal_team.id
                            resource.owner_email = admin_user.email
                            resource.visibility = "public"  # Make visible to all users
                            if hasattr(resource, "federation_source") and not resource.federation_source:
                                resource.federation_source = "mcpgateway-0.7.0-migration"

                        db.commit()
                        total_assigned += len(unassigned)

                except Exception as e:
                    logger.error(f"Failed to assign {resource_name}: {e}")
                    continue

            if total_assigned > 0:
                logger.info(f"Successfully assigned {total_assigned} orphaned resources to admin team")
            else:
                logger.info("No orphaned resources found - all resources have team assignments")

    except Exception as e:
        logger.error(f"Failed to bootstrap resource assignments: {e}")


async def main() -> None:
    """
    Bootstrap or upgrade the database schema, then log readiness.

    Runs `create_all()` + `alembic stamp head` on an empty DB, otherwise just
    executes `alembic upgrade head`, leaving application data intact.
    Also creates the platform admin user if email authentication is enabled.

    Uses distributed advisory locks (PG/MySQL) or file locking (SQLite)
    to prevent race conditions when multiple workers start simultaneously.

    Args:
        None

    Raises:
        Exception: If migration or bootstrap fails
    """
    engine = create_engine(settings.database_url)
    ini_path = files("mcpgateway").joinpath("alembic.ini")
    cfg = Config(str(ini_path))  # path in container
    cfg.attributes["configure_logger"] = True

    # Use advisory lock to prevent concurrent migrations
    try:
        with engine.connect() as conn:
            # Commit any open transaction on the connection before locking (though it should be fresh)
            conn.commit()

            with advisory_lock(conn):
                logger.info("Acquired migration lock, checking database schema...")

                # Pass the LOCKED connection to Alembic config
                cfg.attributes["connection"] = conn

                # Escape '%' characters in URL to avoid configparser interpolation errors
                # (e.g., URL-encoded passwords like %40 for '@')
                escaped_url = settings.database_url.replace("%", "%%")
                cfg.set_main_option("sqlalchemy.url", escaped_url)

                insp = inspect(conn)
                table_names = insp.get_table_names()

                if "gateways" not in table_names:
                    logger.info("Empty DB detected - creating baseline schema")
                    # Apply MariaDB compatibility fixes if needed
                    if settings.database_url.startswith(("mariadb", "mysql")):
                        # pylint: disable=import-outside-toplevel
                        # First-Party
                        from mcpgateway.alembic.env import _modify_metadata_for_mariadb, mariadb_naming_convention

                        _modify_metadata_for_mariadb()
                        Base.metadata.naming_convention = mariadb_naming_convention
                        logger.info("Applied MariaDB compatibility modifications")

                    Base.metadata.create_all(bind=conn)
                    command.stamp(cfg, "head")
                else:
                    versions: list[str] = []
                    if "alembic_version" in table_names:
                        try:
                            rows = conn.execute(text("SELECT version_num FROM alembic_version")).fetchall()
                            versions = [row[0] for row in rows if row[0]]
                        except Exception as exc:
                            logger.warning("Failed to read alembic_version table: %s", exc)

                    if not versions and _schema_looks_current(insp):
                        logger.warning("Existing database has no Alembic revision rows; stamping head to avoid reapplying migrations")
                        command.stamp(cfg, "head")
                    else:
                        logger.info("Running Alembic migrations to ensure schema is up to date")
                        command.upgrade(cfg, "head")

                # Post-upgrade normalization passes (inside lock to be safe)
                updated = normalize_team_visibility(conn)
                if updated:
                    logger.info(f"Normalized {updated} team record(s) to supported visibility values")

                # Bootstrap admin user after database is ready, using the LOCKED connection
                await bootstrap_admin_user(conn)

                # Bootstrap default RBAC roles after admin user is created
                await bootstrap_default_roles(conn)

                # Assign orphaned resources to admin personal team after all setup is complete
                await bootstrap_resource_assignments(conn)

                conn.commit()  # Ensure all migration changes are permanently committed

    except Exception as e:
        logger.error(f"Migration/Bootstrap failed: {e}")
        # Allow retry logic or container restart to handle transient issues
        raise
    finally:
        # Dispose the engine to close all connections in the pool
        engine.dispose()

    logger.info("Database ready")


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