from contextlib import contextmanager
from datetime import datetime, date, time
import logging
import math
import os
from typing import Any

from hypothesis import given
import pytest

from pggm_datalab_utils.db import cursor, query, write, insert_with_return, get_db_connection
from test_utils import sql_friendly_indexed_data, equivalent, sort_by_keys, none_nan_to_zero


@contextmanager
def with_sqlite_db():
    sentinel = object()
    original = os.environ.get('DB_DRIVER', sentinel)
    os.environ['DB_DRIVER'] = 'SQLITE'
    yield
    if original != sentinel: os.environ['DB_DRIVER'] = original


def infer_sqlite_type(column: list[Any]) -> tuple[str, bool]:
    values = [v for v in column if v is not None]
    nullable = len(values) < len(column)

    all_types = {type(value) for value in column if value is not None}
    assert len(all_types) > 0, 'Column is all None'

    if float in all_types and all_types <= {float, int, bool}:  # Float type and all compatible types
        db_type = 'real'
    elif all_types <= {int, bool}:
        db_type = 'integer'
    # Datetime columns
    elif {str} < all_types <= {datetime, date, time, str}:
        db_type = 'datetime'
    elif all_types == {str}:
        db_type = 'text'
    elif list in all_types or dict in all_types:
        db_type = 'json'
    else:
        logging.info(f'Unrecognized type combinations {all_types}, inferring blob type')
        db_type = 'blob'
    return db_type, nullable


def records_to_sqlite_schema(table_name: str, records: list[dict[str, Any]], add_pk: bool = False) -> str:
    assert len(records) > 0
    columns = {k: [r[k] for r in records] for k in records[0].keys()}
    coltypes = {name: infer_sqlite_type(column) for name, column in columns.items()}
    column_sql = ', '.join(f'{name} {typ}' + (' not null' if not nullable else ' null')
                           for name, (typ, nullable) in coltypes.items())
    if add_pk:
        column_sql = 'test_primary_key_autogenerated integer primary key, ' + column_sql

    return f'create table {table_name}({column_sql})'


@given(ipf_data=sql_friendly_indexed_data())
@with_sqlite_db()
def test_read_write_roundtrip(ipf_data):
    ipf_data, index_keys = ipf_data
    ipf_data = list(ipf_data.values())
    ipf = sort_by_keys(ipf_data, index_keys)
    with cursor('localhost', ':memory:') as c:
        c.execute(records_to_sqlite_schema('test', ipf_data))

        write(c, 'test', ipf_data, index_keys, insert=True, update=True, delete=True)
        result = sort_by_keys(query(c, 'select * from test'), index_keys)
    assert equivalent(ipf, result)


@given(ipf_data=sql_friendly_indexed_data(min_size=4))
@with_sqlite_db()
def test_write_no_primary_key(ipf_data):
    ipf_data, index_keys = ipf_data
    ipf_data = list(ipf_data.values())
    ipf = sort_by_keys(ipf_data, index_keys)

    with cursor('localhost', ':memory:') as c:
        c.execute(records_to_sqlite_schema('test', ipf_data))
        with pytest.raises(AssertionError):
            write(c, 'test', ipf_data, insert=True, update=False, delete=True)
        with pytest.raises(AssertionError):
            write(c, 'test', ipf_data, insert=True, update=True, delete=False)
        with pytest.raises(AssertionError):
            write(c, 'test', ipf_data, insert=True, update=True, delete=False)
        write(c, 'test', ipf_data, insert=True, update=False, delete=False)
        result = sort_by_keys(query(c, 'select * from test'), index_keys)
        expected = ipf
        assert equivalent(expected, result)


@given(ipf_data=sql_friendly_indexed_data(min_size=4))
@with_sqlite_db()
def test_primary_key_from_db(ipf_data):
    from copy import deepcopy
    ipf_data, index_keys = ipf_data

    keys = list(ipf_data.keys())
    n = int(len(keys) / 3)

    ipf_original = {k: ipf_data[k] for k in keys[0 * n:2 * n]}
    modified_data = [ipf_data[k] for k in keys[2 * n:3 * n]]

    # Same keys but different values (maybe!)
    ipf_modified = deepcopy(ipf_original)
    for k, new_data in zip(ipf_modified.keys(), modified_data):
        for k_ in ipf_modified[k].keys():
            # Maintain index keys but change non-index keys
            if k_ in index_keys: continue
            ipf_modified[k][k_] = new_data[k_]

    assert len(ipf_modified) == len(ipf_original)
    assert ipf_modified.keys() == ipf_original.keys()
    ipf_modified = list(ipf_modified.values())
    ipf_original = list(ipf_original.values())

    with cursor('localhost', ':memory:') as c:
        c.execute(records_to_sqlite_schema('test', list(ipf_data.values()), add_pk=True))
        # Insert original rows and update keys of modified rows
        for i, row in enumerate(ipf_original):
            ipf_modified[i]['test_primary_key_autogenerated'] = insert_with_return(
                c, 'test', row, 'test_primary_key_autogenerated'
            )['test_primary_key_autogenerated']

        write(c, 'test', ipf_modified, primary_key='test_primary_key_autogenerated',
              insert=True, update=True, delete=True)
        actual = query(c, 'select * from test')
        expected = ipf_modified
        assert equivalent(expected, actual)


@given(ipf_data=sql_friendly_indexed_data(min_size=4))
@with_sqlite_db()
def test_write(ipf_data):
    ipf_data, index_keys = ipf_data

    keys = list(ipf_data.keys())
    n = int(len(keys) / 4)

    ipf_0 = [ipf_data[k] for k in keys[0 * n:1 * n]]
    ipf_1 = [ipf_data[k] for k in keys[1 * n:2 * n]]
    ipf_2 = [ipf_data[k] for k in keys[2 * n:3 * n]]

    # Same keys but different values (maybe!)
    ipf_modified = {k: ipf_data[k_] for k, k_ in zip(keys[1 * n:2 * n], keys[3 * n:4 * n])}
    for key in ipf_modified.keys():
        key_it = key if len(index_keys) > 1 else [key]
        for k, v in zip(index_keys, key_it):
            ipf_modified[key][k] = v
    ipf_modified = list(ipf_modified.values())
    ipf_data = list(ipf_data.values())

    with cursor('localhost', ':memory:') as c:
        c.execute(records_to_sqlite_schema('test', ipf_data))
        write(c, 'test', ipf_0 + ipf_1, index_keys, insert=True, update=True, delete=True)
        write(c, 'test', ipf_1 + ipf_2, index_keys, insert=True, update=False, delete=False)
        result = sort_by_keys(query(c, 'select * from test'), index_keys)
        expected = sort_by_keys(ipf_0 + ipf_1 + ipf_2, index_keys)
        assert equivalent(expected, result)

    # Fresh start
    with cursor('localhost', ':memory:') as c:
        c.execute(records_to_sqlite_schema('test', ipf_data))
        write(c, 'test', ipf_0 + ipf_1, index_keys, insert=True, update=True, delete=True)
        write(c, 'test', ipf_1 + ipf_2, index_keys, insert=True, update=False, delete=True)
        result = sort_by_keys(query(c, 'select * from test'), index_keys)
        expected = sort_by_keys(ipf_1 + ipf_2, index_keys)
        assert equivalent(expected, result)

    with cursor('localhost', ':memory:') as c:
        c.execute(records_to_sqlite_schema('test', ipf_data))
        write(c, 'test', ipf_0 + ipf_1, index_keys, insert=True, update=True, delete=True)
        write(c, 'test', ipf_1 + ipf_2, index_keys, insert=False, update=False, delete=True)
        result = query(c, 'select * from test')
        assert equivalent([], result)

    with cursor('localhost', ':memory:') as c:
        c.execute(records_to_sqlite_schema('test', ipf_data))
        write(c, 'test', ipf_0 + ipf_1, index_keys, insert=True, update=True, delete=True)
        write(c, 'test', ipf_modified + ipf_2, index_keys, insert=False, update=True, delete=False)
        result = sort_by_keys(query(c, 'select * from test'), index_keys)
        expected = sort_by_keys(ipf_0 + ipf_modified, index_keys)
        assert equivalent(expected, result)

    with cursor('localhost', ':memory:') as c:
        c.execute(records_to_sqlite_schema('test', ipf_data))
        write(c, 'test', ipf_0 + ipf_1, index_keys, insert=True, update=True, delete=True)
        write(c, 'test', ipf_modified + ipf_2, index_keys, insert=True, update=True, delete=False)
        result = sort_by_keys(query(c, 'select * from test'), index_keys)
        expected = sort_by_keys(ipf_0 + ipf_modified + ipf_2, index_keys)
        assert equivalent(expected, result)

    with cursor('localhost', ':memory:') as c:
        c.execute(records_to_sqlite_schema('test', ipf_data))
        write(c, 'test', ipf_0 + ipf_1, index_keys, insert=True, update=True, delete=True)
        write(c, 'test', ipf_modified + ipf_2, index_keys, insert=True, update=True, delete=True)
        result = sort_by_keys(query(c, 'select * from test'), index_keys)
        expected = sort_by_keys(ipf_modified + ipf_2, index_keys)
        assert equivalent(expected, result)


# Test a situation where you delete, update, insert a list of dicts into a DB which has two columns of type unique
# And the list of dicts contain a unique primary key, but not a unique value for the second unique column.
# Solution: delete first, then update, then insert.
@with_sqlite_db()
def test_write_two_unique_columns():
    with cursor('localhost', ':memory:') as c:
        c.execute("CREATE TABLE test(ticker type UNIQUE, name, sedol type UNIQUE, benchmark_weight)")
        data_db = [
            ("AAPL", "Apple", "B0YQ5W0", 0.2),
            ("AMZN", "Amazon", "B0YPX78", 0.1)
        ]
        c.executemany("INSERT INTO test VALUES(?, ?, ?, ?)", data_db)
        # insert unique Sedol and unique Ticker
        data_to_insert = [{'ticker': 'AAPL', 'name': 'Apple', 'sedol': 'B0YQ5W0', 'benchmark_weight': 0.2},
                          {'ticker': 'AMZN', 'name': 'Amazon', 'sedol': 'B0YPX78', 'benchmark_weight': 0.1},
                          {'ticker': 'FB', 'name': 'Facebook', 'sedol': 'B7TL820', 'benchmark_weight': 0.4}]
        write(c, 'test', data_to_insert, 'sedol', insert=True, update=True, delete=True)
        # insert company that changed sedol, but not ticker.
        data_to_insert = [{'ticker': 'AAPL', 'name': 'Apple', 'sedol': 'B0YQ5W0', 'benchmark_weight': 0.2},
                          {'ticker': 'AMZN', 'name': 'Amazon', 'sedol': 'B0YPX78', 'benchmark_weight': 0.1},
                          {'ticker': 'FB', 'name': 'Meta', 'sedol': '9KLIKL3', 'benchmark_weight': 0.7}]
        write(c, 'test', data_to_insert, 'sedol', insert=True, update=True, delete=True)
        result = sort_by_keys(query(c, 'select * from test'), ('sedol',))
        expected = sort_by_keys([{'ticker': 'AAPL', 'name': 'Apple', 'sedol': 'B0YQ5W0', 'benchmark_weight': 0.2},
                                 {'ticker': 'AMZN', 'name': 'Amazon', 'sedol': 'B0YPX78', 'benchmark_weight': 0.1},
                                 {'ticker': 'FB', 'name': 'Meta', 'sedol': '9KLIKL3', 'benchmark_weight': 0.7}],
                                ('sedol',))
        assert equivalent(expected, result)


@with_sqlite_db()
def test_only_primary_key():
    data = [{'data': 1}]
    with cursor('localhost', ':memory:') as c:
        c.execute('create table test ( data integer primary key )')
        write(c, 'test', data, primary_key='data', update=True, insert=True, delete=False)
        write(c, 'test', data, primary_key='data', update=True, insert=True, delete=False)
        assert query(c, 'select * from test') == data


@with_sqlite_db()
def test_only_primary_key():
    data = [{'data': 1}]
    with cursor('localhost', ':memory:') as c:
        c.execute('create table test ( data integer primary key )')
        write(c, 'test', data, primary_key='data', update=True, insert=True, delete=False)
        write(c, 'test', data, primary_key='data', update=True, insert=True, delete=False)
        assert query(c, 'select * from test') == data


def test_nan_none():
    data = [{'n': 0.1, 's': None}, {'n': math.nan, 's': 'a'}]
    with cursor('pggm-sql-pilot-lre-o.database.windows.net', 'test') as c:
        c.execute('drop table if exists test')
        # Decimal type
        c.execute('create table test ( n decimal(15, 5), s nvarchar(20) )')
        write(c, 'test', data, primary_key='n', update=True, insert=True, delete=False)
        assert equivalent(query(c, 'select * from test'), data, sort_key=lambda r: none_nan_to_zero(r['n']))

        # Float type
        c.execute('create table test2 ( n real, s nvarchar(20) )')
        write(c, 'test2', data, primary_key='n', update=True, insert=True, delete=False)
        actual = query(c, 'select * from test2')
        assert equivalent(actual, data, sort_key=lambda r: none_nan_to_zero(r['n']))


@pytest.mark.skip('Local test only')
@pytest.mark.parametrize("server, database", [
    ('pggm-sqlmi-ses-1-o.c02c6a164420.database.windows.net', 'FO_MOR'),
    ('SLP-50052', 'VB_UDM_P'),
    ('pggm-sql-cpam-o.database.windows.net', 'cra')
])
def test_db_connection(server, database):
    import pyodbc
    assert isinstance(get_db_connection(server=server, database=database), pyodbc.Connection)
