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()

[((736, 766), 'pytest.fixture', 'pytest.fixture', ([], {'scope': '"""module"""'}), "(scope='module')\n", (750, 766), False, 'import pytest\n'), ((1804, 1854), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""table_name,_"""', 'TEST_DATA'], {}), "('table_name,_', TEST_DATA)\n", (1827, 1854), False, 'import pytest\n'), ((2071, 2129), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""table_name,csv_fpath"""', 'TEST_DATA'], {}), "('table_name,csv_fpath', TEST_DATA)\n", (2094, 2129), False, 'import pytest\n'), ((2531, 2589), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""table_name,csv_fpath"""', 'TEST_DATA'], {}), "('table_name,csv_fpath', TEST_DATA)\n", (2554, 2589), False, 'import pytest\n'), ((2840, 2898), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""table_name,csv_fpath"""', 'TEST_DATA'], {}), "('table_name,csv_fpath', TEST_DATA)\n", (2863, 2898), False, 'import pytest\n'), ((3640, 3698), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""table_name,csv_fpath"""', 'TEST_DATA'], {}), "('table_name,csv_fpath', TEST_DATA)\n", (3663, 3698), False, 'import pytest\n'), ((4000, 4058), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""table_name,csv_fpath"""', 'TEST_DATA'], {}), "('table_name,csv_fpath', TEST_DATA)\n", (4023, 4058), False, 'import pytest\n'), ((4372, 4423), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""table_name,_1"""', 'TEST_DATA'], {}), "('table_name,_1', TEST_DATA)\n", (4395, 4423), False, 'import pytest\n'), ((5748, 5803), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""table_name,_1"""', 'TEST_GIS_DATA'], {}), "('table_name,_1', TEST_GIS_DATA)\n", (5771, 5803), False, 'import pytest\n'), ((6831, 6885), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""ds_name"""', 'TEST_TEMPLATE_DATA'], {}), "('ds_name', TEST_TEMPLATE_DATA)\n", (6854, 6885), False, 'import pytest\n'), ((1160, 1178), 'sqlalchemy.create_engine', 'create_engine', (['uri'], {}), '(uri)\n', (1173, 1178), False, 'from sqlalchemy import create_engine\n'), ((4203, 4223), 'pickle.dumps', 'pickle.dumps', (['source'], {}), '(source)\n', (4215, 4223), False, 'import pickle\n'), ((4243, 4271), 'pickle.loads', 'pickle.loads', (['pickled_source'], {}), '(pickled_source)\n', (4255, 4271), False, 'import pickle\n'), ((4486, 4529), 'os.path.join', 'os.path.join', (['TEST_DATA_DIR', '"""catalog1.yml"""'], {}), "(TEST_DATA_DIR, 'catalog1.yml')\n", (4498, 4529), False, 'import os\n'), ((4545, 4572), 'intake.open_catalog', 'open_catalog', (['catalog_fpath'], {}), '(catalog_fpath)\n', (4557, 4572), False, 'from intake import open_catalog\n'), ((4982, 5024), 'pandas.read_sql_query', 'pd.read_sql_query', (['pgsrc._sql_expr', 'engine'], {}), '(pgsrc._sql_expr, engine)\n', (4999, 5024), True, 'import pandas as pd\n'), ((5153, 5196), 'os.path.join', 'os.path.join', (['TEST_DATA_DIR', '"""catalog1.yml"""'], {}), "(TEST_DATA_DIR, 'catalog1.yml')\n", (5165, 5196), False, 'import os\n'), ((5212, 5239), 'intake.open_catalog', 'open_catalog', (['catalog_fpath'], {}), '(catalog_fpath)\n', (5224, 5239), False, 'from intake import open_catalog\n'), ((5627, 5669), 'pandas.read_sql_query', 'pd.read_sql_query', (['pgsrc._sql_expr', 'engine'], {}), '(pgsrc._sql_expr, engine)\n', (5644, 5669), True, 'import pandas as pd\n'), ((5907, 5950), 'os.path.join', 'os.path.join', (['TEST_DATA_DIR', '"""catalog1.yml"""'], {}), "(TEST_DATA_DIR, 'catalog1.yml')\n", (5919, 5950), False, 'import os\n'), ((5966, 5993), 'intake.open_catalog', 'open_catalog', (['catalog_fpath'], {}), '(catalog_fpath)\n', (5978, 5993), False, 'from intake import open_catalog\n'), ((6375, 6385), 'sqlalchemy.MetaData', 'MetaData', ([], {}), '()\n', (6383, 6385), False, 'from sqlalchemy import MetaData\n'), ((6940, 6983), 'os.path.join', 'os.path.join', (['TEST_DATA_DIR', '"""catalog1.yml"""'], {}), "(TEST_DATA_DIR, 'catalog1.yml')\n", (6952, 6983), False, 'import os\n'), ((6999, 7026), 'intake.open_catalog', 'open_catalog', (['catalog_fpath'], {}), '(catalog_fpath)\n', (7011, 7026), False, 'from intake import open_catalog\n'), ((7390, 7432), 'pandas.read_sql_query', 'pd.read_sql_query', (['pgsrc._sql_expr', 'engine'], {}), '(pgsrc._sql_expr, engine)\n', (7407, 7432), True, 'import pandas as pd\n'), ((1243, 1281), 'os.path.join', 'os.path.join', (['TEST_DATA_DIR', 'csv_fname'], {}), '(TEST_DATA_DIR, csv_fname)\n', (1255, 1281), False, 'import os\n'), ((1295, 1317), 'pandas.read_csv', 'pd.read_csv', (['csv_fpath'], {}), '(csv_fpath)\n', (1306, 1317), True, 'import pandas as pd\n'), ((1439, 1478), 'os.path.join', 'os.path.join', (['TEST_DATA_DIR', 'psql_fname'], {}), '(TEST_DATA_DIR, psql_fname)\n', (1451, 1478), False, 'import os\n'), ((2210, 2248), 'os.path.join', 'os.path.join', (['TEST_DATA_DIR', 'csv_fpath'], {}), '(TEST_DATA_DIR, csv_fpath)\n', (2222, 2248), False, 'import os\n'), ((2666, 2704), 'os.path.join', 'os.path.join', (['TEST_DATA_DIR', 'csv_fpath'], {}), '(TEST_DATA_DIR, csv_fpath)\n', (2678, 2704), False, 'import os\n'), ((3094, 3132), 'os.path.join', 'os.path.join', (['TEST_DATA_DIR', 'csv_fpath'], {}), '(TEST_DATA_DIR, csv_fpath)\n', (3106, 3132), False, 'import os\n'), ((3776, 3814), 'os.path.join', 'os.path.join', (['TEST_DATA_DIR', 'csv_fpath'], {}), '(TEST_DATA_DIR, csv_fpath)\n', (3788, 3814), False, 'import os\n'), ((6617, 6650), 'pandas.read_sql_query', 'pd.read_sql_query', (['_query', 'engine'], {}), '(_query, engine)\n', (6634, 6650), True, 'import pandas as pd\n'), ((6686, 6701), 'shapely.wkt.loads', 'wkt.loads', (['geom'], {}), '(geom)\n', (6695, 6701), False, 'from shapely import wkt\n'), ((6757, 6772), 'shapely.wkt.loads', 'wkt.loads', (['geom'], {}), '(geom)\n', (6766, 6772), False, 'from shapely import wkt\n')]