"""Step definitions for TDD: MEMORY_ENGINES TOCTOU race condition fix.

Tests confirm that MEMORY_ENGINES_LOCK is exported from engine_cache,
that the lock is acquired on engine creation, that cache hits correctly
populate self._engine, and that concurrent access yields a single shared
engine instance.
"""

from __future__ import annotations

import threading
from typing import Any
from behave import given, then, when
from behave.runner import Context


# ---------------------------------------------------------------------------
# Module-level helpers (lazy imports to avoid PYTHONPATH ordering issues)
# ---------------------------------------------------------------------------


def _get_engine_cache_module() -> Any:
    """Return the engine_cache module (lazy import)."""
    import cleveragents.infrastructure.database.engine_cache as m

    return m


def _get_memory_engines() -> Any:
    """Return the live MEMORY_ENGINES cache dictionary."""
    return _get_engine_cache_module().MEMORY_ENGINES


def _get_memory_engines_lock() -> Any:
    """Return the current MEMORY_ENGINES_LOCK from the engine_cache module."""
    return _get_engine_cache_module().MEMORY_ENGINES_LOCK


def _get_uow_class() -> Any:
    """Return the UnitOfWork class (lazy import)."""
    from cleveragents.infrastructure.database.unit_of_work import UnitOfWork

    return UnitOfWork


def _clear_memory_engines() -> None:
    """Remove all entries from the MEMORY_ENGINES cache."""
    # Use the original (real) lock, not any proxy that may be installed.
    engine_cache = _get_engine_cache_module()
    real_lock = getattr(engine_cache, "_original_lock_backup", None)
    if real_lock is None:
        real_lock = engine_cache.MEMORY_ENGINES_LOCK
    with real_lock:
        engine_cache.MEMORY_ENGINES.clear()


def _restore_locks(ctx: Context) -> None:
    """Restore original lock references if the tracking proxy was installed."""
    if hasattr(ctx, "_engine_cache_module") and hasattr(ctx, "_original_lock"):
        setattr(ctx._engine_cache_module, "MEMORY_ENGINES_LOCK", ctx._original_lock)
        if hasattr(ctx._engine_cache_module, "_original_lock_backup"):
            del ctx._engine_cache_module._original_lock_backup
    if hasattr(ctx, "_uow_module") and hasattr(ctx, "_original_uow_lock"):
        setattr(ctx._uow_module, "MEMORY_ENGINES_LOCK", ctx._original_uow_lock)


# ---------------------------------------------------------------------------
# Given
# ---------------------------------------------------------------------------


@given("the engine cache module is imported")
def step_engine_cache_imported(ctx: Context) -> None:
    ctx.engine_cache_module = _get_engine_cache_module()


@given("the engine cache is cleared")
def step_clear_engine_cache(ctx: Context) -> None:
    _clear_memory_engines()


@given('a UnitOfWork is created for "{url}"')
def step_create_first_uow(ctx: Context, url: str) -> None:
    UnitOfWork = _get_uow_class()
    ctx.uow1 = UnitOfWork(url, require_confirmation=False)
    ctx.uow1_engine = ctx.uow1.engine


@given("MEMORY_ENGINES_LOCK acquisition is tracked")
def step_track_lock_acquisition(ctx: Context) -> None:
    """Wrap the module-level lock reference with a counting proxy."""
    count_holder: list[int] = [0]
    ctx.lock_count_holder = count_holder

    engine_cache_module = _get_engine_cache_module()
    real_lock = engine_cache_module.MEMORY_ENGINES_LOCK

    # Stash the real lock so _clear_memory_engines bypasses the proxy.
    engine_cache_module._original_lock_backup = real_lock

    class _CountingLock:
        def acquire(self, blocking: bool = True, timeout: float = -1) -> bool:
            count_holder[0] += 1
            return real_lock.acquire(blocking, timeout)

        def release(self) -> None:
            real_lock.release()

        def __enter__(self) -> _CountingLock:
            count_holder[0] += 1
            real_lock.acquire()
            return self

        def __exit__(self, *args: object) -> None:
            real_lock.release()

    counting_lock = _CountingLock()
    ctx._engine_cache_module = engine_cache_module
    ctx._original_lock = real_lock
    setattr(engine_cache_module, "MEMORY_ENGINES_LOCK", counting_lock)

    import cleveragents.infrastructure.database.unit_of_work as _uow_module

    ctx._uow_module = _uow_module
    ctx._original_uow_lock = _uow_module.MEMORY_ENGINES_LOCK
    setattr(_uow_module, "MEMORY_ENGINES_LOCK", counting_lock)


# ---------------------------------------------------------------------------
# When
# ---------------------------------------------------------------------------


@when('a second UnitOfWork is created for "{url}"')
def step_create_second_uow(ctx: Context, url: str) -> None:
    UnitOfWork = _get_uow_class()
    ctx.uow2 = UnitOfWork(url, require_confirmation=False)
    ctx.uow2_engine = ctx.uow2.engine


@when('a UnitOfWork engine is accessed for "{url}"')
def step_access_uow_engine(ctx: Context, url: str) -> None:
    UnitOfWork = _get_uow_class()
    uow = UnitOfWork(url, require_confirmation=False)
    ctx.accessed_engine = uow.engine
    _restore_locks(ctx)


@when('{count:d} threads each access the engine cache for "{url}"')
def step_concurrent_cache_access(ctx: Context, count: int, url: str) -> None:
    """Directly test the engine cache's lock by calling create_engine in threads.

    This scenario bypasses UnitOfWork migration logic and directly exercises
    the MEMORY_ENGINES check-and-set to verify the lock prevents duplicate creation.
    """
    from sqlalchemy import create_engine

    engines: list[Any] = []
    result_lock = threading.Lock()
    errors: list[Exception] = []
    engine_cache = _get_engine_cache_module()

    def _worker() -> None:
        try:
            # Replicate the check-and-set logic from UnitOfWork.engine.
            with engine_cache.MEMORY_ENGINES_LOCK:
                if url not in engine_cache.MEMORY_ENGINES:
                    engine_cache.MEMORY_ENGINES[url] = create_engine(
                        url,
                        echo=False,
                        future=True,
                        isolation_level="SERIALIZABLE",
                        connect_args={"check_same_thread": False},
                    )
                eng = engine_cache.MEMORY_ENGINES[url]
            with result_lock:
                engines.append(eng)
        except Exception as exc:
            with result_lock:
                errors.append(exc)

    threads = [threading.Thread(target=_worker) for _ in range(count)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    if errors:
        raise AssertionError(f"Worker threads raised exceptions: {errors}") from errors[
            0
        ]

    ctx.concurrent_engines = engines


# ---------------------------------------------------------------------------
# Then
# ---------------------------------------------------------------------------


@then("MEMORY_ENGINES_LOCK should be accessible from the engine_cache module")
def step_lock_accessible(ctx: Context) -> None:
    assert hasattr(ctx.engine_cache_module, "MEMORY_ENGINES_LOCK"), (
        "MEMORY_ENGINES_LOCK not found in engine_cache module"
    )


@then("MEMORY_ENGINES_LOCK should be a threading.Lock instance")
def step_lock_is_threading_lock(ctx: Context) -> None:
    lock = ctx.engine_cache_module.MEMORY_ENGINES_LOCK
    assert hasattr(lock, "acquire") and hasattr(lock, "release"), (
        f"MEMORY_ENGINES_LOCK does not look like a threading.Lock: {type(lock)}"
    )


@then("both UnitOfWork instances should share the same engine instance")
def step_both_share_same_engine(ctx: Context) -> None:
    assert ctx.uow1_engine is ctx.uow2_engine, (
        f"Expected same engine instance for both UnitOfWork objects, "
        f"but got different objects: {ctx.uow1_engine!r} vs {ctx.uow2_engine!r}"
    )


@then("MEMORY_ENGINES_LOCK should have been acquired at least once")
def step_lock_was_acquired(ctx: Context) -> None:
    _restore_locks(ctx)
    count = ctx.lock_count_holder[0]
    assert count >= 1, (
        f"Expected MEMORY_ENGINES_LOCK to be acquired at least once, "
        f"but acquisition count was {count}"
    )


@then("all threads should have received the same engine instance")
def step_all_threads_same_engine(ctx: Context) -> None:
    engines = ctx.concurrent_engines
    assert len(engines) > 0, "No engines were collected from worker threads"
    first = engines[0]
    for i, eng in enumerate(engines[1:], start=1):
        assert eng is first, (
            f"Thread 0 engine ({first!r}) differs from thread {i} engine ({eng!r})"
        )


@then('MEMORY_ENGINES should contain exactly one entry for "{url}"')
def step_cache_has_one_entry(ctx: Context, url: str) -> None:
    lock = _get_memory_engines_lock()
    engines = _get_memory_engines()
    with lock:
        count = sum(1 for k in engines if k == url)
    assert count == 1, (
        f"Expected exactly 1 MEMORY_ENGINES entry for {url!r}, found {count}"
    )
