import datetime

import pytest

import sqlalchemy

from ehrql.backends.base import (

    DefaultSQLBackend,

    MappedTable,

    QueryTable,

    SQLBackend,

    ValidationError,

)

from ehrql.query_engines.base_sql import BaseSQLQueryEngine

from ehrql.query_model.nodes import Column, TableSchema

from ehrql.tables import PatientFrame, Series, table

class BackendFixture(SQLBackend):

    display_name = "Backend Fixture"

    query_engine_class = BaseSQLQueryEngine

    patient_join_column = "PatientId"

    patients = MappedTable(

        source="Patient",

        columns=dict(

            patient_id="PatID",

            date_of_birth="DateOfBirth",

        ),

    )

    events = MappedTable(

        source="events",

        columns=dict(

            date="date",

        ),

    )

    practice_registrations = QueryTable(

        "SELECT patient_id, date_start, date_end FROM some_table"

    )

    @QueryTable.from_function

    def positive_tests(self):

        table_name = self.config.get("table_name", "some_table")

        return f"SELECT patient_id, date FROM {table_name}"

def test_backend_registers_tables():

    """Test that a backend registers its table names"""

    assert set(BackendFixture.tables) == {

        "patients",

        "events",

        "practice_registrations",

        "positive_tests",

    }

def test_mapped_table_sql_with_modified_names():

    table = BackendFixture().get_table_expression(

        "patients",

        TableSchema(

            date_of_birth=Column(datetime.date),

        ),

    )

    sql = str(sqlalchemy.select(table.c.patient_id, table.c.date_of_birth))

    assert sql == 'SELECT "Patient"."PatID", "Patient"."DateOfBirth" \nFROM "Patient"'

def test_mapped_table_sql_with_matching_names():

    table = BackendFixture().get_table_expression(

        "events",

        TableSchema(

            date=Column(datetime.date),

        ),

    )

    sql = str(sqlalchemy.select(table.c.patient_id, table.c.date))

    assert sql == 'SELECT events."PatientId", events.date \nFROM events'

def test_query_table_sql():

    table = BackendFixture().get_table_expression(

        "practice_registrations",

        TableSchema(

            date_start=Column(datetime.date),

            date_end=Column(datetime.date),

        ),

    )

    sql = str(sqlalchemy.select(table.c.patient_id, table.c.date_start))

    assert sql == (

        "SELECT practice_registrations.patient_id, practice_registrations.date_start \n"

        "FROM (SELECT patient_id, date_start, date_end FROM some_table) AS "

        "practice_registrations"

    )

def test_query_table_from_function_sql():

    backend = BackendFixture(config={"table_name": "other_table"})

    table = backend.get_table_expression(

        "positive_tests",

        TableSchema(date=Column(datetime.date)),

    )

    assert str(table) == "SELECT patient_id, date FROM other_table"

def test_default_backend_sql():

    backend = DefaultSQLBackend(query_engine_class=BaseSQLQueryEngine)

    table = backend.get_table_expression(

        "some_table", TableSchema(i=Column(int), b=Column(bool))

    )

    sql = str(sqlalchemy.select(table.c.patient_id, table.c.i, table.c.b))

    assert sql == (

        "SELECT some_table.patient_id, some_table.i, some_table.b \nFROM some_table"

    )

# Use a class as a convenient namespace (`types.SimpleNamespace` would also work)

class Schema:

    @table

    class patients(PatientFrame):

        date_of_birth = Series(datetime.date)

def test_backend_definition_is_allowed_extra_tables_and_columns():

    class BackendFixture(SQLBackend):

        display_name = "Backend Fixture"

        query_engine_class = BaseSQLQueryEngine

        patient_join_column = "patient_id"

        implements = [Schema]

        patients = MappedTable(

            source="patient",

            columns=dict(date_of_birth="DoB", sex="sex"),

        )

        events = MappedTable(

            source="patient",

            columns=dict(date="date", code="code"),

        )

    assert BackendFixture

def test_backend_definition_accepts_query_table():

    class BackendFixture(SQLBackend):

        display_name = "Backend Fixture"

        query_engine_class = BaseSQLQueryEngine

        patient_join_column = "patient_id"

        implements = [Schema]

        patients = QueryTable(

            "SELECT patient_id, CAST(DoB AS date) AS date_of_birth FROM patients",

        )

    assert BackendFixture

def test_backend_definition_fails_if_missing_tables():

    with pytest.raises(ValidationError, match="does not implement table"):

        class BackendFixture(SQLBackend):

            display_name = "Backend Fixture"

            query_engine_class = BaseSQLQueryEngine

            patient_join_column = "patient_id"

            implements = [Schema]

            events = MappedTable(

                source="patient",

                columns=dict(date="date", code="code"),

            )

def test_backend_definition_fails_if_missing_column():

    with pytest.raises(ValidationError, match="missing columns"):

        class BackendFixture(SQLBackend):

            display_name = "Backend Fixture"

            query_engine_class = BaseSQLQueryEngine

            patient_join_column = "patient_id"

            implements = [Schema]

            patients = MappedTable(

                source="patient",

                columns=dict(sex="sex"),

            )

def test_backend_definition_fails_if_query_table_missing_columns():

    with pytest.raises(ValidationError, match="SQL does not reference columns"):

        class BackendFixture(SQLBackend):

            display_name = "Backend Fixture"

            query_engine_class = BaseSQLQueryEngine

            patient_join_column = "patient_id"

            implements = [Schema]

            patients = QueryTable(

                "SELECT patient_id, not_date_of_birth FROM patients",

            )