import asyncio
import faulthandler
import json
import os
import shutil
import sys

# we need to import tmpdir
import tempfile
from collections.abc import AsyncGenerator
from contextlib import suppress
from pathlib import Path
from uuid import UUID, uuid4

import anyio
import orjson
import pytest
from asgi_lifespan import LifespanManager
from blockbuster import blockbuster_ctx
from dotenv import load_dotenv
from fastapi.testclient import TestClient
from httpx import ASGITransport, AsyncClient
from langflow.initial_setup.constants import STARTER_FOLDER_NAME
from langflow.main import create_app
from langflow.services.database.models.api_key.model import ApiKey, UnmaskedApiKeyRead
from langflow.services.database.models.flow.model import Flow, FlowCreate, FlowRead
from langflow.services.database.models.folder.model import Folder
from langflow.services.database.models.transactions.model import TransactionTable
from langflow.services.database.models.user.model import User, UserCreate, UserRead
from langflow.services.database.models.vertex_builds.crud import delete_vertex_builds_by_flow_id_unchecked
from langflow.services.deps import (
    get_auth_service,
    get_db_service,
    get_settings_service,
    is_settings_service_initialized,
    session_scope,
)
from lfx.components.input_output import ChatInput
from lfx.graph import Graph
from lfx.log.logger import logger
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.orm import selectinload
from sqlmodel import Session, SQLModel, create_engine, select
from sqlmodel.ext.asyncio.session import AsyncSession
from sqlmodel.pool import StaticPool
from typer.testing import CliRunner

from tests.api_keys import get_openai_api_key

load_dotenv()


@pytest.fixture(scope="session", autouse=True)
def disable_rate_limiting():
    """Disable rate limiting for all tests to prevent 429 errors during test execution."""
    os.environ["LANGFLOW_RATE_LIMIT_ENABLED"] = "false"
    yield
    os.environ.pop("LANGFLOW_RATE_LIMIT_ENABLED", None)


@pytest.fixture(scope="session", autouse=True)
def disable_models_dev_refresh():
    """Keep the models.dev background refresh out of tests.

    Every app boot otherwise launches a lifespan task that fetches
    https://models.dev/api.json mid-test, which both hits the network and
    trips event-loop-block detectors (pyleak) in whatever test happens to be
    running when the request lands. The bundled static model lists are used
    instead, which is also deterministic.
    """
    os.environ["LANGFLOW_MODELS_DEV_REFRESH"] = "false"
    yield
    os.environ.pop("LANGFLOW_MODELS_DEV_REFRESH", None)


@pytest.fixture(scope="session", autouse=True)
def disable_mcp_auto_init():
    """Keep the MCP server auto-initialization out of tests.

    Every app boot otherwise schedules a lifespan task (``delayed_init_mcp_servers``)
    that, ~10s in, reconciles each project's MCP server config. For apikey/none projects
    that reconciliation spawns ``uvx mcp-proxy`` and makes an outbound connect with no
    bounded timeout, so on a slow/CI runner it hangs until the OS connect timeout (~127s),
    inflating every app-fixture test by ~130s and pushing the heaviest test split past the
    CI step timeout. Skipping it keeps the boot local and deterministic.
    """
    previous_env_value = os.environ.get("LANGFLOW_SKIP_MCP_AUTO_INIT")
    previous_setting = (
        get_settings_service().settings.skip_mcp_auto_init
        if is_settings_service_initialized()
        else (previous_env_value or "").lower() in {"1", "true", "yes", "on"}
    )

    os.environ["LANGFLOW_SKIP_MCP_AUTO_INIT"] = "true"
    if is_settings_service_initialized():
        get_settings_service().set("skip_mcp_auto_init", value=True)
    yield
    if previous_env_value is None:
        os.environ.pop("LANGFLOW_SKIP_MCP_AUTO_INIT", None)
    else:
        os.environ["LANGFLOW_SKIP_MCP_AUTO_INIT"] = previous_env_value
    if is_settings_service_initialized():
        get_settings_service().set("skip_mcp_auto_init", previous_setting)


# TODO: Revert this to True once bb.functions[func].can_block_in("http/client.py", "_safe_read") is fixed
@pytest.fixture(autouse=False)
def blockbuster(request):
    if "benchmark" in request.keywords or "no_blockbuster" in request.keywords:
        yield
    else:
        with blockbuster_ctx() as bb:
            for func in [
                "io.BufferedReader.read",
                "io.BufferedWriter.write",
                "io.TextIOWrapper.read",
                "io.TextIOWrapper.write",
                "os.mkdir",
                "os.stat",
                "os.path.abspath",
            ]:
                bb.functions[func].can_block_in("settings/service.py", "initialize")
            for func in [
                "io.BufferedReader.read",
                "io.TextIOWrapper.read",
            ]:
                bb.functions[func].can_block_in("importlib_metadata/__init__.py", "metadata")
                # bb.functions[func].can_block_in("http/client.py", "_safe_read")

            (
                bb.functions["os.stat"]
                # TODO: make set_class_code async
                .can_block_in("langflow/custom/custom_component/component.py", "set_class_code")
                # TODO: follow discussion in https://github.com/encode/httpx/discussions/3456
                .can_block_in("httpx/_client.py", "_init_transport")
                .can_block_in("rich/traceback.py", "_render_stack")
                .can_block_in("langchain_core/_api/internal.py", "is_caller_internal")
                .can_block_in("langchain_core/runnables/utils.py", "get_function_nonlocals")
                .can_block_in("alembic/versions", "_load_revisions")
                .can_block_in("dotenv/main.py", "find_dotenv")
                .can_block_in("alembic/script/base.py", "_load_revisions")
                .can_block_in("alembic/env.py", "_do_run_migrations")
            )

            for func in ["os.stat", "os.path.abspath", "os.scandir", "os.listdir"]:
                bb.functions[func].can_block_in("alembic/util/pyfiles.py", "load_python_file")
                bb.functions[func].can_block_in("dotenv/main.py", "find_dotenv")
                bb.functions[func].can_block_in("pkgutil.py", "_iter_file_finder_modules")

            for func in ["os.path.abspath", "os.scandir"]:
                bb.functions[func].can_block_in("alembic/script/base.py", "_load_revisions")

            # Add os.stat to alembic/script/base.py _load_revisions
            bb.functions["os.stat"].can_block_in("alembic/script/base.py", "_load_revisions")

            (
                bb.functions["os.path.abspath"]
                .can_block_in("loguru/_better_exceptions.py", {"_get_lib_dirs", "_format_exception"})
                .can_block_in("sqlalchemy/dialects/sqlite/pysqlite.py", "create_connect_args")
                .can_block_in("botocore/__init__.py", "__init__")
            )

            bb.functions["socket.socket.connect"].can_block_in("urllib3/connection.py", "_new_conn")
            bb.functions["ssl.SSLSocket.send"].can_block_in("ssl.py", "sendall")
            bb.functions["ssl.SSLSocket.read"].can_block_in("ssl.py", "recv_into")

            yield bb


# Hard, GIL-proof per-test watchdog. pytest-timeout's thread method (pyproject
# `timeout = 90`) arms a *Python* timer thread, which needs the GIL to run its
# callback -- a hang inside a C call that never releases the GIL (observed on
# the release-1.11.3 py3.13 Group 5 job: a worker froze for 25+ minutes in
# test_login.py::test_session_endpoint_rejects_expired_external_token with no
# dump) silently defeats it and the job burns to the CI step wall.
# faulthandler.dump_traceback_later() instead uses a C-level watchdog thread
# that needs no GIL: it dumps every thread's stack to the real stderr fd
# (inherited by xdist workers, so it lands in the CI log) and, with exit=True,
# hard-exits the wedged worker -- xdist then reports "node down", replaces the
# worker, and --reruns retries the test on the fresh one. The 120s default sits
# above pytest-timeout's 90s so the soft watchdog (clean per-test failure)
# always gets first shot; this only fires when that one *couldn't* run.
# Disable locally for debugger sessions with LANGFLOW_TEST_HARD_TIMEOUT=0.
_HARD_TIMEOUT_S = float(os.getenv("LANGFLOW_TEST_HARD_TIMEOUT", "120"))

# Real-stderr fd, dup'd at pytest_configure time. pytest's fd-level capture
# redirects fd 2 into a per-test temp file that is discarded when the process
# hard-exits, so a dump armed against sys.__stderr__ at *test* time vanishes
# (which is also why pytest-timeout's thread dumps never showed in CI logs).
# At configure time fd 2 still points at the process's original stderr -- in an
# xdist worker that fd is inherited from the controller, so dumps written to
# the dup land in the CI step log. Same strategy as pytest's builtin
# faulthandler plugin.
_watchdog_stderr_fd: int | None = None


@pytest.hookimpl(wrapper=True)
def pytest_runtest_protocol(item, nextitem):  # noqa: ARG001
    if _HARD_TIMEOUT_S <= 0 or _watchdog_stderr_fd is None:
        return (yield)
    faulthandler.dump_traceback_later(_HARD_TIMEOUT_S, file=_watchdog_stderr_fd, exit=True)
    try:
        return (yield)
    finally:
        faulthandler.cancel_dump_traceback_later()


def pytest_configure(config):
    global _watchdog_stderr_fd  # noqa: PLW0603
    if _HARD_TIMEOUT_S > 0 and _watchdog_stderr_fd is None:
        with suppress(AttributeError, ValueError, OSError):
            try:
                fd = sys.stderr.fileno()
            except (AttributeError, ValueError, OSError):
                fd = sys.__stderr__.fileno()
            _watchdog_stderr_fd = os.dup(fd)

    config.addinivalue_line("markers", "noclient: don't create a client for this test")
    config.addinivalue_line("markers", "load_flows: load the flows for this test")
    config.addinivalue_line("markers", "api_key_required: run only if the api key is set in the environment variables")
    config.addinivalue_line(
        "markers",
        "real_services: Tests that need real service instances (real SQLite + real Postgres + real Redis)",
    )
    data_path = Path(__file__).parent.absolute() / "data"

    pytest.BASIC_EXAMPLE_PATH = data_path / "basic_example.json"
    pytest.COMPLEX_EXAMPLE_PATH = data_path / "complex_example.json"
    pytest.OPENAPI_EXAMPLE_PATH = data_path / "Openapi.json"
    pytest.GROUPED_CHAT_EXAMPLE_PATH = data_path / "grouped_chat.json"
    pytest.ONE_GROUPED_CHAT_EXAMPLE_PATH = data_path / "one_group_chat.json"
    pytest.VECTOR_STORE_GROUPED_EXAMPLE_PATH = data_path / "vector_store_grouped.json"
    pytest.WEBHOOK_TEST = data_path / "WebhookTest.json"

    pytest.BASIC_CHAT_WITH_PROMPT_AND_HISTORY = data_path / "BasicChatwithPromptandHistory.json"
    pytest.CHAT_INPUT = data_path / "ChatInputTest.json"
    pytest.TWO_OUTPUTS = data_path / "TwoOutputsTest.json"
    pytest.VECTOR_STORE_PATH = data_path / "Vector_store.json"
    pytest.SIMPLE_API_TEST = data_path / "SimpleAPITest.json"
    pytest.MEMORY_CHATBOT_NO_LLM = data_path / "MemoryChatbotNoLLM.json"
    pytest.ENV_VARIABLE_TEST = data_path / "env_variable_test.json"
    pytest.LOOP_TEST = data_path / "LoopTest.json"
    pytest.CODE_WITH_SYNTAX_ERROR = """
def get_text():
    retun "Hello World"
    """

    # validate that all the paths are correct and the files exist
    for path in [
        pytest.BASIC_EXAMPLE_PATH,
        pytest.COMPLEX_EXAMPLE_PATH,
        pytest.OPENAPI_EXAMPLE_PATH,
        pytest.GROUPED_CHAT_EXAMPLE_PATH,
        pytest.ONE_GROUPED_CHAT_EXAMPLE_PATH,
        pytest.VECTOR_STORE_GROUPED_EXAMPLE_PATH,
        pytest.BASIC_CHAT_WITH_PROMPT_AND_HISTORY,
        pytest.CHAT_INPUT,
        pytest.TWO_OUTPUTS,
        pytest.VECTOR_STORE_PATH,
        pytest.MEMORY_CHATBOT_NO_LLM,
        pytest.LOOP_TEST,
    ]:
        assert path.exists(), f"File {path} does not exist. Available files: {list(data_path.iterdir())}"


# def _has_nonempty_env(var: str) -> bool:
#     return bool((os.getenv(var) or "").strip())


# def pytest_runtest_setup(item):
#     """Auto-skip tests marked with `api_key_required` when no valid OPENAI_API_KEY is provided."""
#     if item.get_closest_marker("api_key_required") and not _has_nonempty_env("OPENAI_API_KEY"):
#         import pytest as _pytest

#         _pytest.skip("OPENAI_API_KEY is not set or is empty")


def pytest_collection_modifyitems(config, items):  # noqa: ARG001
    """Automatically add markers based on test file location."""
    for item in items:
        if "tests/unit/" in str(item.fspath):
            item.add_marker(pytest.mark.unit)
        elif "tests/integration/" in str(item.fspath):
            item.add_marker(pytest.mark.integration)
        elif "tests/slow/" in str(item.fspath):
            item.add_marker(pytest.mark.slow)


async def delete_transactions_by_flow_id(db: AsyncSession, flow_id: UUID):
    if not flow_id:
        return
    stmt = select(TransactionTable).where(TransactionTable.flow_id == flow_id)
    transactions = await db.exec(stmt)
    for transaction in transactions:
        await db.delete(transaction)


async def _delete_transactions_and_vertex_builds(session, flows: list[Flow]):
    from langflow.services.database.models.jobs.model import Job

    flow_ids = [flow.id for flow in flows]
    for flow_id in flow_ids:
        if not flow_id:
            continue
        try:
            # Delete associated jobs first due to foreign keys
            stmt = select(Job).where(Job.flow_id == flow_id)
            jobs = (await session.exec(stmt)).all()
            for job in jobs:
                await session.delete(job)
            await session.flush()

            await delete_vertex_builds_by_flow_id_unchecked(session, flow_id)
        except Exception as e:
            logger.debug(f"Error deleting jobs/vertex builds for flow {flow_id}: {e}")
        try:
            await delete_transactions_by_flow_id(session, flow_id)
        except Exception as e:
            logger.debug(f"Error deleting transactions for flow {flow_id}: {e}")


@pytest.fixture
async def async_client() -> AsyncGenerator:
    app = create_app()
    async with (
        LifespanManager(app, startup_timeout=None, shutdown_timeout=60) as manager,
        AsyncClient(transport=ASGITransport(app=manager.app), base_url="http://testserver", http2=True) as client,
    ):
        yield client


@pytest.fixture(name="session")
def session_fixture():
    engine = create_engine(
        "sqlite+pysqlite:///:memory:",
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )
    # Ensure foreign keys are enabled for the in-memory session engine
    from sqlalchemy import event

    @event.listens_for(engine, "connect")
    def set_sqlite_pragma(dbapi_connection, connection_record):  # noqa: ARG001
        cursor = dbapi_connection.cursor()
        cursor.execute("PRAGMA foreign_keys=ON")
        cursor.close()

    try:
        SQLModel.metadata.create_all(engine)
        with Session(engine) as session:
            yield session
    finally:
        SQLModel.metadata.drop_all(engine)
        engine.dispose()


@pytest.fixture
async def async_session():
    engine = create_async_engine("sqlite+aiosqlite://", connect_args={"check_same_thread": False}, poolclass=StaticPool)
    try:
        async with engine.begin() as conn:
            await conn.run_sync(SQLModel.metadata.create_all)
        async with AsyncSession(engine, expire_on_commit=False) as session:
            yield session
        async with engine.begin() as conn:
            await conn.run_sync(SQLModel.metadata.drop_all)
    finally:
        await engine.dispose()


class Config:
    broker_url = "redis://localhost:6379/0"
    result_backend = "redis://localhost:6379/0"


@pytest.fixture(name="load_flows_dir")
def load_flows_dir():
    with tempfile.TemporaryDirectory() as tempdir:
        yield tempdir


@pytest.fixture(name="distributed_env")
def _setup_env(monkeypatch):
    monkeypatch.setenv("LANGFLOW_CACHE_TYPE", "redis")
    monkeypatch.setenv("LANGFLOW_REDIS_HOST", "result_backend")
    monkeypatch.setenv("LANGFLOW_REDIS_PORT", "6379")
    monkeypatch.setenv("LANGFLOW_REDIS_DB", "0")
    monkeypatch.setenv("LANGFLOW_REDIS_EXPIRE", "3600")
    monkeypatch.setenv("LANGFLOW_REDIS_PASSWORD", "")
    monkeypatch.setenv("FLOWER_UNAUTHENTICATED_API", "True")
    monkeypatch.setenv("BROKER_URL", "redis://result_backend:6379/0")
    monkeypatch.setenv("RESULT_BACKEND", "redis://result_backend:6379/0")
    monkeypatch.setenv("C_FORCE_ROOT", "true")


@pytest.fixture(name="distributed_client")
def distributed_client_fixture(
    session: Session,  # noqa: ARG001
    monkeypatch,
    distributed_env,  # noqa: ARG001
):
    # Here we load the .env from ../deploy/.env
    from langflow.core import celery_app

    db_dir = tempfile.mkdtemp()
    try:
        db_path = Path(db_dir) / "test.db"
        monkeypatch.setenv("LANGFLOW_DATABASE_URL", f"sqlite:///{db_path}")
        monkeypatch.setenv("LANGFLOW_AUTO_LOGIN", "false")
        monkeypatch.setenv("LANGFLOW_SUPERUSER", "langflow")
        monkeypatch.setenv("LANGFLOW_SUPERUSER_PASSWORD", "test-superuser-password")
        # monkeypatch langflow.services.task.manager.USE_CELERY to True
        # monkeypatch.setattr(manager, "USE_CELERY", True)
        monkeypatch.setattr(celery_app, "celery_app", celery_app.make_celery("langflow", Config))

        # def get_session_override():
        #     return session

        app = create_app()

        # app.dependency_overrides[get_session] = get_session_override
        with TestClient(app) as client:
            yield client
    finally:
        shutil.rmtree(db_dir)  # Clean up the temporary directory
    app.dependency_overrides.clear()
    monkeypatch.undo()


def get_graph(type_="basic"):
    """Get a graph from a json file."""
    if type_ == "basic":
        path = pytest.BASIC_EXAMPLE_PATH
    elif type_ == "complex":
        path = pytest.COMPLEX_EXAMPLE_PATH
    elif type_ == "openapi":
        path = pytest.OPENAPI_EXAMPLE_PATH

    with path.open(encoding="utf-8") as f:
        flow_graph = json.load(f)
    data_graph = flow_graph["data"]
    nodes = data_graph["nodes"]
    edges = data_graph["edges"]
    graph = Graph()
    graph.add_nodes_and_edges(nodes, edges)
    return graph


@pytest.fixture
def basic_graph_data():
    with pytest.BASIC_EXAMPLE_PATH.open(encoding="utf-8") as f:
        return json.load(f)


@pytest.fixture
def basic_graph():
    return get_graph()


@pytest.fixture
def complex_graph():
    return get_graph("complex")


@pytest.fixture
def openapi_graph():
    return get_graph("openapi")


@pytest.fixture
def json_flow():
    return pytest.BASIC_EXAMPLE_PATH.read_text(encoding="utf-8")


@pytest.fixture
def grouped_chat_json_flow():
    return pytest.GROUPED_CHAT_EXAMPLE_PATH.read_text(encoding="utf-8")


@pytest.fixture
def one_grouped_chat_json_flow():
    return pytest.ONE_GROUPED_CHAT_EXAMPLE_PATH.read_text(encoding="utf-8")


@pytest.fixture
def vector_store_grouped_json_flow():
    return pytest.VECTOR_STORE_GROUPED_EXAMPLE_PATH.read_text(encoding="utf-8")


@pytest.fixture
def json_flow_with_prompt_and_history():
    return pytest.BASIC_CHAT_WITH_PROMPT_AND_HISTORY.read_text(encoding="utf-8")


@pytest.fixture
def json_simple_api_test():
    return pytest.SIMPLE_API_TEST.read_text(encoding="utf-8")


@pytest.fixture
def json_vector_store():
    return pytest.VECTOR_STORE_PATH.read_text(encoding="utf-8")


@pytest.fixture
def json_webhook_test():
    return pytest.WEBHOOK_TEST.read_text(encoding="utf-8")


@pytest.fixture
def json_memory_chatbot_no_llm():
    return pytest.MEMORY_CHATBOT_NO_LLM.read_text(encoding="utf-8")


@pytest.fixture
def json_loop_test():
    return pytest.LOOP_TEST.read_text(encoding="utf-8")


@pytest.fixture(autouse=True)
def deactivate_tracing(monkeypatch):
    monkeypatch.setenv("LANGFLOW_DEACTIVATE_TRACING", "true")
    yield
    monkeypatch.undo()


@pytest.fixture(autouse=True)
def disable_telemetry_writer(monkeypatch):
    # Tests assert on freshly-written transactions / vertex_builds rows. The
    # batched writer is a production optimization; in tests we want the
    # synchronous legacy DB path so reads-after-writes are visible.
    monkeypatch.setenv("LANGFLOW_TELEMETRY_WRITER_ENABLED", "false")
    yield
    monkeypatch.undo()


@pytest.fixture
def use_noop_session(monkeypatch):
    monkeypatch.setenv("LANGFLOW_USE_NOOP_DATABASE", "1")
    # Optionally patch the Settings object if needed
    # from lfx.services.settings.base import Settings
    # monkeypatch.setattr(Settings, "use_noop_database", True)
    yield
    monkeypatch.undo()


@pytest.fixture(name="client")
async def client_fixture(
    session: Session,  # noqa: ARG001
    monkeypatch,
    request,
    load_flows_dir,
):
    # Set the database url to a test database
    if "noclient" in request.keywords:
        yield
    else:

        def init_app():
            db_dir = tempfile.mkdtemp()
            db_path = Path(db_dir) / "test.db"
            monkeypatch.setenv("LANGFLOW_DATABASE_URL", f"sqlite:///{db_path}")
            monkeypatch.setenv("LANGFLOW_AUTO_LOGIN", "false")
            monkeypatch.setenv("LANGFLOW_SUPERUSER", "langflow")
            monkeypatch.setenv("LANGFLOW_SUPERUSER_PASSWORD", "test-superuser-password")
            monkeypatch.setenv("DO_NOT_TRACK", "true")
            if "load_flows" in request.keywords:
                shutil.copyfile(
                    pytest.BASIC_EXAMPLE_PATH, Path(load_flows_dir) / "c54f9130-f2fa-4a3e-b22a-3856d946351b.json"
                )
                monkeypatch.setenv("LANGFLOW_LOAD_FLOWS_PATH", load_flows_dir)
                monkeypatch.setenv("LANGFLOW_AUTO_LOGIN", "true")
            # Clear the services cache
            from lfx.services.manager import get_service_manager

            get_service_manager().factories.clear()
            get_service_manager().services.clear()  # Clear the services cache
            app = create_app()
            db_service = get_db_service()
            db_service.database_url = f"sqlite:///{db_path}"
            db_service.reload_engine()
            return app, db_path

        app, db_path = await asyncio.to_thread(init_app)
        # app.dependency_overrides[get_session] = get_session_override
        async with (
            LifespanManager(app, startup_timeout=None, shutdown_timeout=60) as manager,
            AsyncClient(transport=ASGITransport(app=manager.app), base_url="http://testserver/", http2=True) as client,
        ):
            yield client
        # app.dependency_overrides.clear()
        monkeypatch.undo()
        # clear the temp db
        with suppress(FileNotFoundError):
            await anyio.Path(db_path).unlink()


@pytest.fixture
def runner(tmp_path):
    env = {"LANGFLOW_DATABASE_URL": f"sqlite:///{tmp_path}/test.db"}
    return CliRunner(env=env)


@pytest.fixture
async def test_user(client):
    user_data = UserCreate(
        username="testuser",
        password="testpassword",  # noqa: S106
    )
    response = await client.post("api/v1/users/", json=user_data.model_dump())
    assert response.status_code == 201
    user = response.json()
    yield user
    # Clean up
    await client.delete(f"/api/v1/users/{user['id']}")


@pytest.fixture
async def active_user(client):  # noqa: ARG001
    async with session_scope() as session:
        user = User(
            username="activeuser",
            password=get_auth_service().get_password_hash("testpassword"),
            is_active=True,
            is_superuser=False,
        )
        stmt = select(User).where(User.username == user.username)
        if active_user := (await session.exec(stmt)).first():
            user = active_user
        else:
            session.add(user)
            await session.flush()
            await session.refresh(user)
        user = UserRead.model_validate(user, from_attributes=True)
    yield user
    # Clean up
    # Now cleanup transactions, vertex_build
    try:
        async with session_scope() as session:
            user = await session.get(User, user.id, options=[selectinload(User.flows)])
            await _delete_transactions_and_vertex_builds(session, user.flows)
    except Exception as e:
        logger.exception(f"Error deleting transactions and vertex builds for user: {e}")

    try:
        async with session_scope() as session:
            user = await session.get(User, user.id)
            await session.delete(user)
    except Exception as e:
        logger.exception(f"Error deleting user: {e}")


@pytest.fixture
async def logged_in_headers(client, active_user):
    login_data = {"username": active_user.username, "password": "testpassword"}
    response = await client.post("api/v1/login", data=login_data)
    assert response.status_code == 200
    tokens = response.json()
    a_token = tokens["access_token"]
    return {"Authorization": f"Bearer {a_token}"}


@pytest.fixture
async def active_super_user(client):  # noqa: ARG001
    async with session_scope() as session:
        user = User(
            username="activeuser",
            password=get_auth_service().get_password_hash("testpassword"),
            is_active=True,
            is_superuser=True,
        )
        stmt = select(User).where(User.username == user.username)
        if active_user := (await session.exec(stmt)).first():
            user = active_user
        else:
            session.add(user)
            await session.flush()
            await session.refresh(user)
        user = UserRead.model_validate(user, from_attributes=True)
    yield user
    # Clean up
    # Now cleanup transactions, vertex_build
    async with session_scope() as session:
        user = await session.get(User, user.id, options=[selectinload(User.flows)])
        await _delete_transactions_and_vertex_builds(session, user.flows)
        await session.delete(user)


@pytest.fixture
async def logged_in_headers_super_user(client, active_super_user):
    login_data = {"username": active_super_user.username, "password": "testpassword"}
    response = await client.post("api/v1/login", data=login_data)
    assert response.status_code == 200
    tokens = response.json()
    a_token = tokens["access_token"]
    return {"Authorization": f"Bearer {a_token}"}


@pytest.fixture
async def flow(
    client,  # noqa: ARG001
    json_flow: str,
    active_user,
):
    loaded_json = json.loads(json_flow)
    flow_data = FlowCreate(name="test_flow", data=loaded_json.get("data"), user_id=active_user.id)

    flow = Flow.model_validate(flow_data.model_dump(exclude={"id"}))
    async with session_scope() as session:
        session.add(flow)
        await session.flush()
        await session.refresh(flow)
        flow = FlowRead.model_validate(flow, from_attributes=True)

    yield flow

    async with session_scope() as session:
        db_flow = await session.get(Flow, flow.id)
        if db_flow:
            await session.delete(db_flow)


@pytest.fixture
def json_chat_input():
    return pytest.CHAT_INPUT.read_text(encoding="utf-8")


@pytest.fixture
def json_two_outputs():
    return pytest.TWO_OUTPUTS.read_text(encoding="utf-8")


@pytest.fixture
async def added_flow_webhook_test(client, json_webhook_test, logged_in_headers):
    flow = orjson.loads(json_webhook_test)
    data = flow["data"]
    flow = FlowCreate(name="Basic Chat", description="description", data=data)
    response = await client.post("api/v1/flows/", json=flow.model_dump(), headers=logged_in_headers)
    assert response.status_code == 201
    assert response.json()["name"] == flow.name
    assert response.json()["data"] == flow.data
    yield response.json()
    await client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)


@pytest.fixture
async def added_flow_chat_input(client, json_chat_input, logged_in_headers):
    flow = orjson.loads(json_chat_input)
    data = flow["data"]
    flow = FlowCreate(name="Chat Input", description="description", data=data)
    response = await client.post("api/v1/flows/", json=flow.model_dump(), headers=logged_in_headers)
    assert response.status_code == 201
    assert response.json()["name"] == flow.name
    assert response.json()["data"] == flow.data
    yield response.json()
    await client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)


@pytest.fixture
async def added_flow_two_outputs(client, json_two_outputs, logged_in_headers):
    flow = orjson.loads(json_two_outputs)
    data = flow["data"]
    flow = FlowCreate(name="Two Outputs", description="description", data=data)
    response = await client.post("api/v1/flows/", json=flow.model_dump(), headers=logged_in_headers)
    assert response.status_code == 201
    assert response.json()["name"] == flow.name
    assert response.json()["data"] == flow.data
    yield response.json()
    await client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)


@pytest.fixture
async def added_vector_store(client, json_vector_store, logged_in_headers):
    vector_store = orjson.loads(json_vector_store)
    data = vector_store["data"]
    vector_store = FlowCreate(name="Vector Store", description="description", data=data)
    response = await client.post("api/v1/flows/", json=vector_store.model_dump(), headers=logged_in_headers)
    assert response.status_code == 201
    assert response.json()["name"] == vector_store.name
    assert response.json()["data"] == vector_store.data
    yield response.json()
    await client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)


@pytest.fixture
async def added_webhook_test(client, json_webhook_test, logged_in_headers):
    webhook_test = orjson.loads(json_webhook_test)
    data = webhook_test["data"]
    webhook_test = FlowCreate(
        name="Webhook Test", description="description", data=data, endpoint_name=webhook_test["endpoint_name"]
    )
    response = await client.post("api/v1/flows/", json=webhook_test.model_dump(), headers=logged_in_headers)
    assert response.status_code == 201
    assert response.json()["name"] == webhook_test.name
    assert response.json()["data"] == webhook_test.data
    yield response.json()
    await client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)


@pytest.fixture
async def flow_component(client: AsyncClient, logged_in_headers):
    chat_input = ChatInput()
    graph = Graph(start=chat_input, end=chat_input)
    graph_dict = graph.dump(name="Chat Input Component")
    flow = FlowCreate(**graph_dict)
    response = await client.post("api/v1/flows/", json=flow.model_dump(), headers=logged_in_headers)
    assert response.status_code == 201
    yield response.json()
    await client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)


@pytest.fixture
async def created_api_key(active_user):
    hashed = get_auth_service().get_password_hash("random_key")
    api_key = ApiKey(
        name="test_api_key",
        user_id=active_user.id,
        api_key="random_key",
        hashed_api_key=hashed,
    )
    async with session_scope() as session:
        stmt = select(ApiKey).where(ApiKey.api_key == api_key.api_key)
        if existing_api_key := (await session.exec(stmt)).first():
            existing_api_key = UnmaskedApiKeyRead.model_validate(existing_api_key, from_attributes=True)
            yield existing_api_key
            return
        session.add(api_key)
        await session.flush()
        await session.refresh(api_key)
        api_key = UnmaskedApiKeyRead.model_validate(api_key, from_attributes=True)

    yield api_key

    async with session_scope() as session:
        db_key = await session.get(ApiKey, api_key.id)
        if db_key:
            await session.delete(db_key)


@pytest.fixture
def user_one_api_key(created_api_key: ApiKey) -> str:
    """Provides the API key for user_one."""
    return created_api_key.api_key


@pytest.fixture
async def user_two(
    client: AsyncClient,  # noqa: ARG001
) -> AsyncGenerator[User, None]:
    """Creates a second user for multi-user access tests."""
    user_id = uuid4()
    async with session_scope() as session:
        user = User(
            id=user_id,
            username=f"test_user_two_{user_id}",
            password=get_auth_service().get_password_hash("hashed_password"),
            is_active=True,
        )
        session.add(user)
        await session.flush()
        await session.refresh(user)
        user = UserRead.model_validate(user, from_attributes=True)

    yield user

    async with session_scope() as session:
        # Cleanup related API keys first
        keys_to_delete = (await session.exec(select(ApiKey).where(ApiKey.user_id == user.id))).all()
        for key in keys_to_delete:
            await session.delete(key)

        # Cleanup the user
        user_to_delete = await session.get(User, user.id)
        if user_to_delete:
            await session.delete(user_to_delete)


@pytest.fixture
async def created_user_two_api_key(user_two: User) -> AsyncGenerator[ApiKey, None]:
    """Creates and yields an API key for the second user."""
    raw_key = f"user-two-key-{uuid4()}"
    hashed_key = get_auth_service().get_password_hash(raw_key)
    api_key = ApiKey(
        user_id=user_two.id,
        name="Test API Key for User Two",
        api_key=raw_key,
        hashed_api_key=hashed_key,
    )

    async with session_scope() as session:
        session.add(api_key)
        await session.flush()
        await session.refresh(api_key)
        api_key = UnmaskedApiKeyRead.model_validate(api_key, from_attributes=True)

    yield api_key

    async with session_scope() as session:
        db_key = await session.get(ApiKey, api_key.id)
        if db_key:
            await session.delete(db_key)


@pytest.fixture
def user_two_api_key(created_user_two_api_key: ApiKey) -> str:
    """Provides the API key string for the second user."""
    return created_user_two_api_key.api_key


@pytest.fixture(name="simple_api_test")
async def get_simple_api_test(client, logged_in_headers, json_simple_api_test):
    # Once the client is created, we can get the starter project
    # Just create a new flow with the simple api test
    flow = orjson.loads(json_simple_api_test)
    data = flow["data"]
    flow = FlowCreate(name="Simple API Test", data=data, description="Simple API Test")
    response = await client.post("api/v1/flows/", json=flow.model_dump(), headers=logged_in_headers)
    assert response.status_code == 201
    yield response.json()
    await client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)


@pytest.fixture(name="starter_project")
async def get_starter_project(client, active_user):  # noqa: ARG001
    # once the client is created, we can get the starter project
    async with session_scope() as session:
        stmt = (
            select(Flow)
            .where(Flow.folder.has(Folder.name == STARTER_FOLDER_NAME))
            .where(Flow.name == "Basic Prompting")
        )
        flow = (await session.exec(stmt)).first()
        if not flow:
            msg = "No starter project found"
            raise ValueError(msg)

        # ensure openai api key is set
        openai_api_key = get_openai_api_key()
        data_as_json = json.dumps(flow.data)
        data_as_json = data_as_json.replace("OPENAI_API_KEY", openai_api_key)
        # also replace `"load_from_db": true` with `"load_from_db": false`
        if '"load_from_db": true' in data_as_json:
            data_as_json = data_as_json.replace('"load_from_db": true', '"load_from_db": false')
        if '"load_from_db": true' in data_as_json:
            msg = "load_from_db should be false"
            raise ValueError(msg)
        flow.data = json.loads(data_as_json)

        new_flow_create = FlowCreate(
            name=flow.name,
            description=flow.description,
            data=flow.data,
            user_id=active_user.id,
        )
        new_flow = Flow.model_validate(new_flow_create, from_attributes=True)
        session.add(new_flow)
        await session.flush()
        await session.refresh(new_flow)
        new_flow_dict = new_flow.model_dump()
        yield new_flow_dict
        # Clean up
        await session.delete(new_flow)