import datetime
import re

import pytest

from ehrql import Dataset


@pytest.fixture(params=["execute", "dump_sql"])
def spec_test(request, engine):
    # Test that we can insert the data, run the query, and get the expected results
    def run_test_execute(table_data, series, expected_results, population=None):
        # Populate database tables.
        engine.populate(
            {
                table: parse_table(table._qm_node.schema, s)
                for table, s in table_data.items()
            }
        )

        # Create a Dataset with the specified population and a single variable which is
        # the series under test.
        dataset = make_dataset(table_data, population)
        dataset.v = series

        # If we're comparing floats then we want only approximate equality to account
        # for rounding differences
        if series._type is float:
            expected_results = pytest.approx(expected_results, rel=1e-5)

        # Extract data, and check it's as expected.
        results = [(r["patient_id"], r["v"]) for r in engine.extract(dataset)]
        results_dict = dict(results)
        assert len(results) == len(results_dict), "Duplicate patient IDs found"
        assert results_dict == expected_results

        # Assert types are as expected
        for patient_id, value in results_dict.items():
            if value is not None:
                assert isinstance(value, series._type), (
                    f"Expected {series._type} got {type(value)} in "
                    f"result {{{patient_id}: {value}}}"
                )

    # Test that we can generate SQL with literal parmeters for debugging purposes
    def run_test_dump_sql(table_data, series, expected_results, population=None):
        # Create a Dataset with the specified population and a single variable which is
        # the series under test.
        dataset = make_dataset(table_data, population)
        dataset.v = series

        # Check that we can generate SQL without error
        assert engine.dump_dataset_sql(dataset)

    mode = request.param

    if mode == "execute":
        return run_test_execute
    elif mode == "dump_sql":
        if engine.name == "in_memory":
            pytest.skip("in_memory engine produces no SQL")
        return run_test_dump_sql
    else:
        assert False


def make_dataset(table_data, population=None):
    # To reduce noise in the tests we provide a default population which contains all
    # patients in any tables referenced in the data
    if population is None:
        population = False
        for table in table_data.keys():
            population = table.exists_for_patient() | population
    dataset = Dataset()
    dataset.define_population(population)
    return dataset


def parse_table(schema, s):
    """Parse string containing table data, returning list of dicts.

    See test_conftest.py for examples.
    """

    header, _, *lines = s.strip().splitlines()
    col_names = [token.strip() for token in header.split("|")]
    col_names[0] = "patient_id"
    column_types = dict(
        patient_id=int, **{name: type_ for name, type_ in schema.column_types}
    )
    rows = [parse_row(column_types, col_names, line) for line in lines]
    return rows


def parse_row(column_types, col_names, line):
    """Parse string containing row data, returning list of values.

    See test_conftest.py for examples.
    """

    # Regex splits on any '|' character, as long as it's not adjacent
    # to another '|' character using look-ahead and look-behind. This
    # is to allow '||' to appear as content within a field, currently
    # just for the all_diagnoses and all_procedures fields in apcs
    return {
        col_name: parse_value(column_types[col_name], token.strip())
        for col_name, token in zip(col_names, re.split(r"(?<!\|)\|(?!\|)", line))
    }


def parse_value(type_, value):
    """Parse string returning value of correct type for column.

    An empty string indicates a null value.
    """
    if not value:
        return None

    if hasattr(type_, "_primitive_type"):
        type_ = type_._primitive_type()

    if type_ is bool:
        parse = lambda v: {"T": True, "F": False}[v]  # noqa E731
    elif type_ == datetime.date:
        parse = datetime.date.fromisoformat
    else:
        parse = type_

    return parse(value)