import os
import signal
import subprocess
import sys
import time
from contextlib import contextmanager

from backend.app_paths import get_db_schema, get_db_url, get_workload_folder_path, is_remote_storage
from backend.execution import Execution
from backend.sqlite_compat import SQLITE_DRIVER
from backend.utils.logging_utils import get_logger

logger = get_logger(__name__)


def _format_subprocess_error(process_name: str, returncode: int) -> str:
    if returncode == -signal.SIGKILL:
        return f"{process_name} was killed by the OS (SIGKILL). " "This usually indicates a memory issue on the host."
    return f"{process_name} failed"


@contextmanager
def timed_step(step_name: str):
    """Context manager to time and log a step."""
    logger.info("[%s] Starting...", step_name)
    start = time.perf_counter()
    try:
        yield
        duration = time.perf_counter() - start
        logger.info("[%s] Completed in %.2fs", step_name, duration)
    except Exception as e:
        duration = time.perf_counter() - start
        logger.error("[%s] Failed after %.2fs: %s", step_name, duration, e)
        raise


def _log_app_tables(engine, schema, table_log_label: str):
    """Log all app tables from SQLAlchemy metadata after migration."""
    try:
        from backend.database.base import db
        from sqlalchemy import inspect

        inspector = inspect(engine)
        existing_tables = set(inspector.get_table_names(schema=schema))

        # Get expected tables from SQLAlchemy metadata (source of truth)
        # Import models to ensure metadata is populated
        from backend.database import models  # noqa: F401

        expected_tables = sorted(db.metadata.tables.keys())

        logger.info("%s (schema=%s):", table_log_label, schema)
        for table in expected_tables:
            status = "✓" if table in existing_tables else "✗"
            logger.info("  %s %s", status, table)
    except Exception as e:
        logger.warning("Could not list app tables: %s", e)


def _setup_alembic_env(
    db_url: str, workload_folder_path: str | None, schema: str | None, tables_prefix: str, app_mode: str
):
    env = os.environ.copy()
    env.update(
        {
            "DATABASE_URL": db_url,
            "APP_MODE": app_mode,
        }
    )
    if workload_folder_path:
        env["DB_FOLDER_PATH"] = workload_folder_path
    if schema:
        env["DB_SCHEMA"] = schema
    if tables_prefix:
        env["TABLES_PREFIX"] = tables_prefix
    plugin_python_lib = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../python-lib"))
    logger.info("Plugin python-lib path: %s", plugin_python_lib)
    # TODO SQLAlchemy/Pydantic: Check if this is correct in all cases
    logger.info("Using executable: %s", sys.executable)
    resource_folder = os.environ.get("DKU_CUSTOM_RESOURCE_FOLDER")
    if resource_folder:
        project_root = os.path.dirname(resource_folder)
    else:
        project_root = Execution("").exec_path
    logger.info("Project root for Alembic: %s", project_root)

    return sys.executable, env, project_root


def run_alembic_downgrade(
    db_url: str, workload_folder_path: str | None, schema: str | None, tables_prefix: str, app_mode: str
):
    python_executable, env, project_root = _setup_alembic_env(
        db_url, workload_folder_path, schema, tables_prefix, app_mode
    )
    logger.info("Running Alembic downgrade")

    try:
        logger.info("Downgrading the database to the base revision")
        subprocess.run([python_executable, "-m", "alembic", "downgrade", "-1"], check=True, env=env, cwd=project_root)
    except Exception as e:
        logger.exception(f"Failed to run Alembic downgrade: {e}")
        raise RuntimeError("Alembic downgrade failed") from e


def run_alembic_upgrade(
    db_url: str,
    workload_folder_path: str | None,
    schema: str | None,
    tables_prefix: str,
    app_mode: str,
    table_log_label: str,
):
    python_executable, env, project_root = _setup_alembic_env(
        db_url, workload_folder_path, schema, tables_prefix, app_mode
    )
    try:
        from sqlalchemy import create_engine

        engine = create_engine(db_url)
        logger.info("Starting Alembic upgrade for DB: %s", db_url)

        result = subprocess.run(
            [python_executable, "-m", "alembic", "upgrade", "head"],
            check=True,
            env=env,
            cwd=project_root,
            capture_output=True,
            text=True,
        )

        if result.stdout:
            for line in result.stdout.strip().split("\n"):
                if line:
                    logger.info("%s", line)
        if result.stderr:
            for line in result.stderr.strip().split("\n"):
                if line:
                    logger.info("%s", line)

        if "Running upgrade" in result.stdout or "Running upgrade" in result.stderr:
            logger.info("Migration applied successfully.")
        else:
            logger.info("Database already at latest revision.")

        _log_app_tables(engine, schema, table_log_label)
        return result
    except subprocess.CalledProcessError as e:
        logger.error("Alembic migration failed with exit code %s", e.returncode)
        logger.error(f"STDOUT: {e.stdout}")
        logger.error(f"STDERR: {e.stderr}")
        raise RuntimeError(_format_subprocess_error("Alembic migration", e.returncode)) from e
    except Exception as e:
        logger.exception(f"Failed to run Alembic migration: {e}")
        raise RuntimeError("Alembic migration failed") from e


def bootstrap_backend(app, app_mode: str, tables_prefix: str, table_log_label: str) -> None:
    schema = get_db_schema()
    if schema:
        os.environ["DB_SCHEMA"] = schema

    if tables_prefix:
        os.environ["TABLES_PREFIX"] = tables_prefix

    db_url = get_db_url()
    workload_folder_path = None if is_remote_storage() else get_workload_folder_path()

    logger.info("========== Backend Startup ==========")
    logger.info("SQLite driver selected: %s", SQLITE_DRIVER)
    total_start = time.perf_counter()

    with timed_step("Database migration"):
        run_alembic_upgrade(db_url, workload_folder_path, schema, tables_prefix, app_mode, table_log_label)

    with timed_step("App setup"):
        # Import after TABLES_PREFIX is configured so models pick up the correct prefix.
        from backend.setup import setup_app

        logger.info("calling setup_app in %s mode", app_mode)
        setup_app(app, db_url, workload_folder_path, app_mode=app_mode)

    logger.info("========== Startup complete in %.2fs ==========", time.perf_counter() - total_start)
