import os
import pickle
import pytest
import pandas as pd
from shapely import wkt

from intake_postgres import PostgresSource
from intake import open_catalog
from .util import verify_datasource_interface


TEST_DATA_DIR = 'tests'
TEST_DATA = [
    ('sample1', 'sample1.csv'),
    ('sample2_1', 'sample2_1.csv'),
    ('sample2_2', 'sample2_2.csv'),
]
TEST_GIS_DATA = [
    ('points', 'sample_points.psql'),
    ('multipoints', 'sample_multipoints.psql'),
    ('lines', 'sample_lines.psql'),
    ('multilines', 'sample_multilines.psql'),
    ('polygons', 'sample_polygons.psql'),
    ('multipolygons', 'sample_multipolygons.psql'),
    # ('triangles', 'sample_triangles.psql'),
]
TEST_TEMPLATE_DATA = [
    'jinja2_params_with_env',
]


@pytest.fixture(scope='module')
def engine():
    """Start docker container for PostgreSQL database, yield a tuple (engine,
    metadata), and cleanup connection afterward."""
    from .util import start_postgres, stop_postgres
    from sqlalchemy import create_engine
    stop_postgres(let_fail=True)
    local_port = start_postgres()

    uri = 'postgresql://postgres@localhost:{}/postgres'.format(local_port)
    engine = create_engine(uri)
    for table_name, csv_fname in TEST_DATA:
        csv_fpath = os.path.join(TEST_DATA_DIR, csv_fname)
        df = pd.read_csv(csv_fpath)
        df.to_sql(table_name, engine, index=False)
    for table_name, psql_fname in TEST_GIS_DATA:
        psql_fpath = os.path.join(TEST_DATA_DIR, psql_fname)
        with engine.connect() as conn:
            with open(psql_fpath, 'r') as fp:
                cmds = fp.read().strip().split(';')
            for cmd in cmds:
                if cmd.strip():
                    conn.execute(' '.join(cmd.split()))

    try:
        yield engine
    finally:
        stop_postgres()


@pytest.mark.parametrize('table_name,_', TEST_DATA)
def test_open(engine, table_name, _):
    d = PostgresSource(str(engine.url), 'select * from '+table_name)
    assert d.container == 'dataframe'
    assert d.description is None
    verify_datasource_interface(d)


@pytest.mark.parametrize('table_name,csv_fpath', TEST_DATA)
def test_discover(engine, table_name, csv_fpath):
    expected_df = pd.read_csv(os.path.join(TEST_DATA_DIR, csv_fpath))
    source = PostgresSource(str(engine.url), 'select * from '+table_name)
    info = source.discover()
    dt = {k: str(v) for k, v in expected_df.dtypes.to_dict().items()}
    assert info['dtype'] == dt
    assert info['shape'] == (None, 3)
    assert info['npartitions'] == 1


@pytest.mark.parametrize('table_name,csv_fpath', TEST_DATA)
def test_read(engine, table_name, csv_fpath):
    expected_df = pd.read_csv(os.path.join(TEST_DATA_DIR, csv_fpath))
    source = PostgresSource(str(engine.url), 'select * from '+table_name)
    df = source.read()
    assert expected_df.equals(df)


@pytest.mark.parametrize('table_name,csv_fpath', TEST_DATA)
def test_discover_after_read(engine, table_name, csv_fpath):
    """Assert that after reading the dataframe, discover() shows more accurate
    information.
    """
    expected_df = pd.read_csv(os.path.join(TEST_DATA_DIR, csv_fpath))
    source = PostgresSource(str(engine.url), 'select * from '+table_name)
    info = source.discover()
    dt = {k: str(v) for k, v in expected_df.dtypes.to_dict().items()}
    assert info['dtype'] == dt
    assert info['shape'] == (None, 3)
    assert info['npartitions'] == 1

    df = source.read()
    assert expected_df.equals(df)

    info = source.discover()
    assert info['dtype'] == dt
    assert info['shape'] == (4, 3)
    assert info['npartitions'] == 1

    assert expected_df.equals(df)


@pytest.mark.parametrize('table_name,csv_fpath', TEST_DATA)
def test_close(engine, table_name, csv_fpath):
    expected_df = pd.read_csv(os.path.join(TEST_DATA_DIR, csv_fpath))
    source = PostgresSource(str(engine.url), 'select * from '+table_name)

    source.close()
    # Can reopen after close
    df = source.read()

    assert expected_df.equals(df)


@pytest.mark.parametrize('table_name,csv_fpath', TEST_DATA)
def test_pickle(engine, table_name, csv_fpath):
    source = PostgresSource(str(engine.url), 'select * from '+table_name)

    pickled_source = pickle.dumps(source)
    source_clone = pickle.loads(pickled_source)

    expected_df = source.read()
    df = source_clone.read()

    assert expected_df.equals(df)


@pytest.mark.parametrize('table_name,_1', TEST_DATA)
def test_catalog(engine, table_name, _1):
    catalog_fpath = os.path.join(TEST_DATA_DIR, 'catalog1.yml')

    catalog = open_catalog(catalog_fpath)
    ds_name = table_name.rsplit('_idx', 1)[0]
    src = catalog[ds_name]
    pgsrc = src.get()
    pgsrc._uri = str(engine.url)

    assert src.describe()['container'] == 'dataframe'
    assert src.describe_open()['plugin'] == 'postgres'
    assert src.describe_open()['args']['sql_expr'][:6] in ('select', 'SELECT')

    metadata = pgsrc.discover()
    assert metadata['npartitions'] == 1

    expected_df = pd.read_sql_query(pgsrc._sql_expr, engine)
    df = pgsrc.read()
    assert expected_df.equals(df)

    pgsrc.close()


def test_catalog_join(engine):
    catalog_fpath = os.path.join(TEST_DATA_DIR, 'catalog1.yml')

    catalog = open_catalog(catalog_fpath)
    ds_name = 'sample2'
    src = catalog[ds_name]
    pgsrc = src.get()
    pgsrc._uri = str(engine.url)

    assert src.describe()['container'] == 'dataframe'
    assert src.describe_open()['plugin'] == 'postgres'
    assert src.describe_open()['args']['sql_expr'][:6] in ('select', 'SELECT')

    metadata = pgsrc.discover()
    assert metadata['npartitions'] == 1

    expected_df = pd.read_sql_query(pgsrc._sql_expr, engine)
    df = pgsrc.read()
    assert expected_df.equals(df)

    pgsrc.close()


@pytest.mark.parametrize('table_name,_1', TEST_GIS_DATA)
def test_postgis_data(engine, table_name, _1):
    from sqlalchemy import MetaData
    catalog_fpath = os.path.join(TEST_DATA_DIR, 'catalog1.yml')

    catalog = open_catalog(catalog_fpath)
    ds_name = table_name
    src = catalog[ds_name]
    pgsrc = src.get()
    pgsrc._uri = str(engine.url)

    assert src.describe()['container'] == 'dataframe'
    assert src.describe_open()['plugin'] == 'postgres'
    assert src.describe_open()['args']['sql_expr'][:6] in ('select', 'SELECT')

    metadata = pgsrc.discover()
    assert metadata['npartitions'] == 1

    meta = MetaData()
    meta.reflect(bind=engine)
    col_exprs = ['ST_AsText({0}) as {0}'.format(col.name)
                 for col in meta.tables[table_name].columns]
    _query = pgsrc._sql_expr.replace('*', ', '.join(col_exprs))
    expected_df = pd.read_sql_query(_query, engine).applymap(
        lambda geom: str(wkt.loads(geom))
    )
    df = pgsrc.read().applymap(lambda geom: str(wkt.loads(geom)))
    assert expected_df.equals(df)

    pgsrc.close()


@pytest.mark.parametrize('ds_name', TEST_TEMPLATE_DATA)
def test_jinja2(engine, ds_name):
    catalog_fpath = os.path.join(TEST_DATA_DIR, 'catalog1.yml')

    catalog = open_catalog(catalog_fpath)
    src = catalog[ds_name]
    pgsrc = src.get()
    pgsrc._uri = str(engine.url)

    assert src.describe()['container'] == 'dataframe'
    assert src.describe_open()['plugin'] == 'postgres'
    assert src.describe_open()['args']['sql_expr'][:6] in ('select', 'SELECT')

    metadata = pgsrc.discover()
    assert metadata['npartitions'] == 1

    expected_df = pd.read_sql_query(pgsrc._sql_expr, engine)
    df = pgsrc.read()
    assert expected_df.equals(df)

    pgsrc.close()

[((32, 1, 32, 31), 'pytest.fixture', 'pytest.fixture', (), '', False, 'import pytest\n'), ((62, 1, 62, 51), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(62, 25, 62, 39): '"""table_name,_"""', (62, 41, 62, 50): 'TEST_DATA'}, {}), "('table_name,_', TEST_DATA)", False, 'import pytest\n'), ((70, 1, 70, 59), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(70, 25, 70, 47): '"""table_name,csv_fpath"""', (70, 49, 70, 58): 'TEST_DATA'}, {}), "('table_name,csv_fpath', TEST_DATA)", False, 'import pytest\n'), ((81, 1, 81, 59), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(81, 25, 81, 47): '"""table_name,csv_fpath"""', (81, 49, 81, 58): 'TEST_DATA'}, {}), "('table_name,csv_fpath', TEST_DATA)", False, 'import pytest\n'), ((89, 1, 89, 59), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(89, 25, 89, 47): '"""table_name,csv_fpath"""', (89, 49, 89, 58): 'TEST_DATA'}, {}), "('table_name,csv_fpath', TEST_DATA)", False, 'import pytest\n'), ((113, 1, 113, 59), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(113, 25, 113, 47): '"""table_name,csv_fpath"""', (113, 49, 113, 58): 'TEST_DATA'}, {}), "('table_name,csv_fpath', TEST_DATA)", False, 'import pytest\n'), ((125, 1, 125, 59), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(125, 25, 125, 47): '"""table_name,csv_fpath"""', (125, 49, 125, 58): 'TEST_DATA'}, {}), "('table_name,csv_fpath', TEST_DATA)", False, 'import pytest\n'), ((138, 1, 138, 52), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(138, 25, 138, 40): '"""table_name,_1"""', (138, 42, 138, 51): 'TEST_DATA'}, {}), "('table_name,_1', TEST_DATA)", False, 'import pytest\n'), ((185, 1, 185, 56), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(185, 25, 185, 40): '"""table_name,_1"""', (185, 42, 185, 55): 'TEST_GIS_DATA'}, {}), "('table_name,_1', TEST_GIS_DATA)", False, 'import pytest\n'), ((217, 1, 217, 55), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(217, 25, 217, 34): '"""ds_name"""', (217, 36, 217, 54): 'TEST_TEMPLATE_DATA'}, {}), "('ds_name', TEST_TEMPLATE_DATA)", False, 'import pytest\n'), ((42, 13, 42, 31), 'sqlalchemy.create_engine', 'create_engine', ({(42, 27, 42, 30): 'uri'}, {}), '(uri)', False, 'from sqlalchemy import create_engine\n'), ((129, 21, 129, 41), 'pickle.dumps', 'pickle.dumps', ({(129, 34, 129, 40): 'source'}, {}), '(source)', False, 'import pickle\n'), ((130, 19, 130, 47), 'pickle.loads', 'pickle.loads', ({(130, 32, 130, 46): 'pickled_source'}, {}), '(pickled_source)', False, 'import pickle\n'), ((140, 20, 140, 63), 'os.path.join', 'os.path.join', ({(140, 33, 140, 46): 'TEST_DATA_DIR', (140, 48, 140, 62): '"""catalog1.yml"""'}, {}), "(TEST_DATA_DIR, 'catalog1.yml')", False, 'import os\n'), ((142, 14, 142, 41), 'intake.open_catalog', 'open_catalog', ({(142, 27, 142, 40): 'catalog_fpath'}, {}), '(catalog_fpath)', False, 'from intake import open_catalog\n'), ((155, 18, 155, 60), 'pandas.read_sql_query', 'pd.read_sql_query', ({(155, 36, 155, 51): 'pgsrc._sql_expr', (155, 53, 155, 59): 'engine'}, {}), '(pgsrc._sql_expr, engine)', True, 'import pandas as pd\n'), ((163, 20, 163, 63), 'os.path.join', 'os.path.join', ({(163, 33, 163, 46): 'TEST_DATA_DIR', (163, 48, 163, 62): '"""catalog1.yml"""'}, {}), "(TEST_DATA_DIR, 'catalog1.yml')", False, 'import os\n'), ((165, 14, 165, 41), 'intake.open_catalog', 'open_catalog', ({(165, 27, 165, 40): 'catalog_fpath'}, {}), '(catalog_fpath)', False, 'from intake import open_catalog\n'), ((178, 18, 178, 60), 'pandas.read_sql_query', 'pd.read_sql_query', ({(178, 36, 178, 51): 'pgsrc._sql_expr', (178, 53, 178, 59): 'engine'}, {}), '(pgsrc._sql_expr, engine)', True, 'import pandas as pd\n'), ((188, 20, 188, 63), 'os.path.join', 'os.path.join', ({(188, 33, 188, 46): 'TEST_DATA_DIR', (188, 48, 188, 62): '"""catalog1.yml"""'}, {}), "(TEST_DATA_DIR, 'catalog1.yml')", False, 'import os\n'), ((190, 14, 190, 41), 'intake.open_catalog', 'open_catalog', ({(190, 27, 190, 40): 'catalog_fpath'}, {}), '(catalog_fpath)', False, 'from intake import open_catalog\n'), ((203, 11, 203, 21), 'sqlalchemy.MetaData', 'MetaData', ({}, {}), '()', False, 'from sqlalchemy import MetaData\n'), ((219, 20, 219, 63), 'os.path.join', 'os.path.join', ({(219, 33, 219, 46): 'TEST_DATA_DIR', (219, 48, 219, 62): '"""catalog1.yml"""'}, {}), "(TEST_DATA_DIR, 'catalog1.yml')", False, 'import os\n'), ((221, 14, 221, 41), 'intake.open_catalog', 'open_catalog', ({(221, 27, 221, 40): 'catalog_fpath'}, {}), '(catalog_fpath)', False, 'from intake import open_catalog\n'), ((233, 18, 233, 60), 'pandas.read_sql_query', 'pd.read_sql_query', ({(233, 36, 233, 51): 'pgsrc._sql_expr', (233, 53, 233, 59): 'engine'}, {}), '(pgsrc._sql_expr, engine)', True, 'import pandas as pd\n'), ((44, 20, 44, 58), 'os.path.join', 'os.path.join', ({(44, 33, 44, 46): 'TEST_DATA_DIR', (44, 48, 44, 57): 'csv_fname'}, {}), '(TEST_DATA_DIR, csv_fname)', False, 'import os\n'), ((45, 13, 45, 35), 'pandas.read_csv', 'pd.read_csv', ({(45, 25, 45, 34): 'csv_fpath'}, {}), '(csv_fpath)', True, 'import pandas as pd\n'), ((48, 21, 48, 60), 'os.path.join', 'os.path.join', ({(48, 34, 48, 47): 'TEST_DATA_DIR', (48, 49, 48, 59): 'psql_fname'}, {}), '(TEST_DATA_DIR, psql_fname)', False, 'import os\n'), ((72, 30, 72, 68), 'os.path.join', 'os.path.join', ({(72, 43, 72, 56): 'TEST_DATA_DIR', (72, 58, 72, 67): 'csv_fpath'}, {}), '(TEST_DATA_DIR, csv_fpath)', False, 'import os\n'), ((83, 30, 83, 68), 'os.path.join', 'os.path.join', ({(83, 43, 83, 56): 'TEST_DATA_DIR', (83, 58, 83, 67): 'csv_fpath'}, {}), '(TEST_DATA_DIR, csv_fpath)', False, 'import os\n'), ((94, 30, 94, 68), 'os.path.join', 'os.path.join', ({(94, 43, 94, 56): 'TEST_DATA_DIR', (94, 58, 94, 67): 'csv_fpath'}, {}), '(TEST_DATA_DIR, csv_fpath)', False, 'import os\n'), ((115, 30, 115, 68), 'os.path.join', 'os.path.join', ({(115, 43, 115, 56): 'TEST_DATA_DIR', (115, 58, 115, 67): 'csv_fpath'}, {}), '(TEST_DATA_DIR, csv_fpath)', False, 'import os\n'), ((208, 18, 208, 51), 'pandas.read_sql_query', 'pd.read_sql_query', ({(208, 36, 208, 42): '_query', (208, 44, 208, 50): 'engine'}, {}), '(_query, engine)', True, 'import pandas as pd\n'), ((209, 25, 209, 40), 'shapely.wkt.loads', 'wkt.loads', ({(209, 35, 209, 39): 'geom'}, {}), '(geom)', False, 'from shapely import wkt\n'), ((211, 48, 211, 63), 'shapely.wkt.loads', 'wkt.loads', ({(211, 58, 211, 62): 'geom'}, {}), '(geom)', False, 'from shapely import wkt\n')]