import os
import random
import subprocess
import threading
from pathlib import Path

import pytest
from hypothesis.internal.reflection import extract_lambda_source

import ehrql
import ehrql.__main__
from ehrql import query_language as ql
from ehrql.main import get_sql_strings
from ehrql.query_engines.in_memory import InMemoryQueryEngine
from ehrql.query_engines.mssql import MSSQLQueryEngine
from ehrql.query_engines.sqlite import SQLiteQueryEngine
from ehrql.query_engines.trino import TrinoQueryEngine
from ehrql.query_model import nodes as qm
from tests.lib.orm_utils import make_orm_models

from .lib.databases import (
    InMemoryPythonDatabase,
    InMemorySQLiteDatabase,
    make_mssql_database,
    make_trino_database,
    wait_for_database,
)
from .lib.docker import Containers


def pytest_collection_modifyitems(session, config, items):  # pragma: no cover
    """If running with pytest-xdist, add a group identifier to each test item, based on
    which database is used by the test.

    This lets us use pytest-xdist to distribute tests across three processes leading to
    a moderate speed-up, via `pytest -n3`.

    The "proper" way to distribute tests with pytest-xdist is by adding the xdist_group
    mark.  However, this is very hard to do dynamically (because of our use of
    request.getfixturevalue) so it is less invasive to add a group identifier here,
    during test collection.  Later, pytest-xdist will use the group identifier to
    distribute tests to workers.
    """

    if "PYTEST_XDIST_WORKER" not in os.environ:
        # Modifying test item identifiers makes it harder to copy and paste identifiers
        # from failing outputs, so it only makes sense to do so if we're running tests
        # with pytest-xdist.
        return

    slow_database_names = ["mssql", "trino"]

    for item in items:
        group = "other"

        if "engine" in item.fixturenames:
            database_name = item.callspec.params["engine"]
            if database_name in slow_database_names:
                group = database_name

        else:
            found_database_in_fixtures = False
            for database_name in slow_database_names:
                if any(
                    database_name in fixture_name for fixture_name in item.fixturenames
                ):
                    group = database_name
                    # Check that tests do not use multiple fixtures for slow databases.
                    assert not found_database_in_fixtures
                    found_database_in_fixtures = True

        item._nodeid = f"{item.nodeid}@{group}"


# Fail the build if we see any warnings.
def pytest_terminal_summary(terminalreporter, exitstatus, config):
    if terminalreporter.stats.get("warnings"):  # pragma: no cover
        print("ERROR: warnings detected")
        if terminalreporter._session.exitstatus == 0:
            terminalreporter._session.exitstatus = 13


def pytest_make_parametrize_id(config, val):
    # Where we use lambdas as test parameters, having the source as the parameter ID
    # makes it quicker to identify specific test cases in the output
    if callable(val) and val.__name__ == "<lambda>":
        return extract_lambda_source(val).removeprefix("lambda: ")


@pytest.fixture(scope="session")
def show_delayed_warning(request):
    """
    Some fixtures can take a long time to execute the first time they're run (e.g. they
    might need to pull down a large Docker image) but pytest's output capturing means
    that the user has no idea what's happening. This fixture allows us to "poke through"
    the output capturing and display a message to the user, but only if the task has
    already taken more than N seconds.
    """

    def show_warning(message):  # pragma: no cover
        capturemanager = request.config.pluginmanager.getplugin("capturemanager")
        # No need to display anything if output is not being captured
        if capturemanager.is_capturing():
            with capturemanager.global_and_fixture_disabled():
                print(f"\n => {message} ...")

    return lambda delay, message: ContextTimer(delay, show_warning, args=[message])


# Timer which starts/cancels itself when entering/exiting a context block
class ContextTimer(threading.Timer):
    def __enter__(self):
        self.start()

    def __exit__(self, *_):
        self.cancel()


@pytest.fixture(scope="session")
def containers():
    yield Containers()


# Database fixtures {

# These fixtures come in pairs.  For each database, there is a session-scoped fixture,
# which performs any setup, and there is a function-scoped fixture, which reuses the
# fixture returned by the session-scoped fixture.
#
# In most cases, we will want the function-scoped fixture, as this allows post-test
# teardown.  However, the generative tests require a session-scoped fixture.


@pytest.fixture(scope="session")
def in_memory_sqlite_database_with_session_scope():
    return InMemorySQLiteDatabase()


@pytest.fixture(scope="function")
def in_memory_sqlite_database(in_memory_sqlite_database_with_session_scope):
    database = in_memory_sqlite_database_with_session_scope
    yield database
    database.teardown()


@pytest.fixture(scope="session")
def mssql_database_with_session_scope(containers, show_delayed_warning):
    with show_delayed_warning(
        3, "Starting MSSQL Docker image (will download image on first run)"
    ):
        database = make_mssql_database(containers)
        wait_for_database(database)
    return database


@pytest.fixture(scope="function")
def mssql_database(mssql_database_with_session_scope):
    database = mssql_database_with_session_scope
    yield database
    database.teardown()


@pytest.fixture(scope="session")
def trino_database_with_session_scope(containers, show_delayed_warning):
    with show_delayed_warning(
        3, "Starting Trino Docker image (will download image on first run)"
    ):
        database = make_trino_database(containers)
        wait_for_database(database)
    return database


@pytest.fixture(scope="function")
def trino_database(trino_database_with_session_scope):
    database = trino_database_with_session_scope
    yield database
    database.teardown()


class QueryEngineFixture:
    def __init__(self, name, database, query_engine_class):
        self.name = name
        self.database = database
        self.query_engine_class = query_engine_class

    def setup(self, *items, metadata=None):
        return self.database.setup(*items, metadata=metadata)

    def teardown(self):
        return self.database.teardown()

    def populate(self, *args):
        return self.setup(make_orm_models(*args))

    def query_engine(self, dsn=False, **engine_kwargs):
        if dsn is False:
            dsn = self.database.host_url()
        return self.query_engine_class(dsn, **engine_kwargs)

    def get_results_tables(self, dataset, **engine_kwargs):
        if isinstance(dataset, ql.Dataset):
            dataset = dataset._compile()
        assert isinstance(dataset, qm.Dataset)
        query_engine = self.query_engine(**engine_kwargs)
        results_tables = query_engine.get_results_tables(dataset)
        # We don't explicitly order the results and not all databases naturally
        # return in the same order
        return [
            [row._asdict() for row in sort_table(table)] for table in results_tables
        ]

    def extract(self, dataset, **engine_kwargs):
        return self.get_results_tables(dataset, **engine_kwargs)[0]

    def dump_dataset_sql(self, dataset, **engine_kwargs):
        assert isinstance(dataset, ql.Dataset)
        dataset_qm = dataset._compile()
        query_engine = self.query_engine(dsn=None, **engine_kwargs)
        return get_sql_strings(query_engine, dataset_qm)

    def sqlalchemy_engine(self):
        return self.query_engine().engine


def sort_table(table):
    # Python won't naturally compare None with other values, but we need to sort tables
    # containg None values so we treat None as smaller than all other values
    return sorted(table, key=sort_key_with_nones)


def sort_key_with_nones(row):
    return [(v is not None, v) for v in row]


QUERY_ENGINE_NAMES = ("in_memory", "sqlite", "mssql", "trino")


def engine_factory(request, engine_name, with_session_scope=False):
    if engine_name == "in_memory":
        return QueryEngineFixture(
            engine_name, InMemoryPythonDatabase(), InMemoryQueryEngine
        )

    if engine_name == "sqlite":
        database_fixture_name = "in_memory_sqlite_database"
        query_engine_class = SQLiteQueryEngine
    elif engine_name == "mssql":
        database_fixture_name = "mssql_database"
        query_engine_class = MSSQLQueryEngine
    elif engine_name == "trino":
        database_fixture_name = "trino_database"
        query_engine_class = TrinoQueryEngine
    else:
        assert False

    if with_session_scope:
        database_fixture_name = f"{database_fixture_name}_with_session_scope"

    # We dynamically request fixtures rather than making them arguments in the usual way
    # so that we only start the database containers we actually need for the test run
    database = request.getfixturevalue(database_fixture_name)

    return QueryEngineFixture(engine_name, database, query_engine_class)


@pytest.fixture(params=QUERY_ENGINE_NAMES)
def engine(request):
    return engine_factory(request, request.param)


@pytest.fixture
def mssql_engine(request):
    return engine_factory(request, "mssql")


@pytest.fixture
def trino_engine(request):
    return engine_factory(request, "trino")


@pytest.fixture
def in_memory_engine(request):
    return engine_factory(request, "in_memory")


@pytest.fixture
def sqlite_engine(request):
    return engine_factory(request, "sqlite")


@pytest.fixture(scope="session")
def ehrql_image(show_delayed_warning):
    project_dir = Path(ehrql.__file__).parents[1]
    # Note different name from production image to avoid confusion
    image = "ehrql-dev"
    # We're deliberately choosing to shell out to the docker client here rather than use
    # the docker-py library to avoid possible difference in the build process (docker-py
    # doesn't seem to be particularly actively maintained)
    with show_delayed_warning(3, f"Building {image} Docker image"):
        subprocess.run(
            ["docker", "build", project_dir, "-t", image],
            check=True,
            env=dict(os.environ, DOCKER_BUILDKIT="1"),
        )
    return f"{image}:latest"


@pytest.fixture(autouse=True)
def random_should_not_be_used():
    """Asserts that every test should leave the global random number generator unchanged.

    This is because we want all of our use of randomness to be based on seeded random number
    generators, so it should not depend on the global random number generator in any way.
    If this is failing, please find all uses of the random module and replace them with a
    Random instance, ideally from a predictable seed.
    """
    prev_state = random.getstate()
    yield
    assert random.getstate() == prev_state, (
        "Global random number generator was used in test."
    )


@pytest.fixture
def call_cli(capsys):
    """
    Wrapper around the CLI entrypoint to make it easier to call from tests
    """

    def call(*args, environ=None):
        # Convert any Path instances to strings
        args = [str(arg) if isinstance(arg, Path) else arg for arg in args]
        ehrql.__main__.main(args, environ=environ)
        return capsys.readouterr()

    # Allow reading captured output even when call throws an exception
    call.readouterr = capsys.readouterr

    return call


@pytest.fixture
def call_cli_docker(containers, ehrql_image):
    """
    As above, but invoke the CLI via the Docker image
    """

    def call(*args, environ=None, workspace=None):
        args = [
            # Make any paths relative to the workspace directory so they still point to
            # the right place inside Docker. If you supply path arguments and no
            # workspace this will error, as it should. Likewise if you supply paths
            # outside of the workspace.
            str(arg.relative_to(workspace)) if isinstance(arg, Path) else str(arg)
            for arg in args
        ]
        if workspace is not None:
            # Because the files in these directories will need to be readable by
            # low-privilege, isolated processes we can't use the standard restrictive
            # permissions for temporary directories
            workspace.chmod(0o755)
            volumes = {workspace: {"bind": "/workspace", "mode": "rw"}}
        else:
            volumes = {}
        return containers.run_captured(
            ehrql_image,
            command=args,
            volumes=volumes,
            environment=environ or {},
        )

    return call