import datetime as dt
import os

from pathlib import Path

import pandas as pd
import pytest

from hvplot.tests.util import makeMixedDataFrame

from lumen.sources.base import Source
from lumen.state import state
from lumen.transforms.sql import SQLLimit


@pytest.fixture
def source(make_filesource):
    root = os.path.dirname(__file__)
    return make_filesource(root)

@pytest.fixture
def expected_df(mixed_df, column_value_type):
    df = mixed_df
    column, value, type = column_value_type

    if type == 'single_value':
        return df[df[column] == value]

    elif type == 'range':
        begin, end = value
        return df[(df[column] >= begin) & (df[column] <= end)]

    elif type == 'range_list':
        conditions = False
        for range in value:
            begin, end = range
            conditions |= ((df[column] >= begin) & (df[column] <= end))
        return df[conditions]

    elif type == 'list':
        return df[df[column].isin(value)]

    elif type == 'date':
        return df[df[column] == pd.to_datetime(value)]

    elif type == 'date_range':
        begin, end = value
        return df[(df[column] >= pd.to_datetime(begin)) & (df[column] <= pd.to_datetime(end))]

    return df


def test_source_resolve_module_type():
    assert Source._get_type('lumen.sources.base.Source') is Source


@pytest.mark.parametrize("filter_col_A", [314, (13, 314), ['A', 314, 'def']])
@pytest.mark.parametrize("filter_col_B", [(3, 15.9)])
@pytest.mark.parametrize("filter_col_C", [[1, 'A', 'def']])
@pytest.mark.parametrize("sql_transforms", [(None, None), (SQLLimit(limit=100), SQLLimit(limit=100))])
def test_file_source_table_cache_key(source, filter_col_A, filter_col_B, filter_col_C, sql_transforms):
    t1, t2 = sql_transforms
    kwargs1 = {}
    kwargs2 = {}
    if t1 is not None:
        kwargs1['sql_transforms'] = [t1]
    if t2 is not None:
        kwargs2['sql_transforms'] = [t2]
    key1 = source._get_key('test', A=filter_col_A, B=filter_col_B, C=filter_col_C, **kwargs1)
    key2 = source._get_key('test', A=filter_col_A, B=filter_col_B, C=filter_col_C, **kwargs2)
    assert key1 == key2


@pytest.mark.parametrize(
    "column_value_type", [
        ('A', 1, 'single_value'),
        ('A', (1, 3), 'range'),
        ('A', (1, 2), 'range'),
        ('A', [(0, 1), (3, 4)], 'range_list'),
        ('C', 'foo2', 'single_value'),
        ('C', ['foo1', 'foo3'], 'list'),
        ('D', dt.datetime(2009, 1, 2), 'single_value'),
        ('D', (dt.datetime(2009, 1, 2), dt.datetime(2009, 1, 5)), 'range'),
        ('D', [dt.datetime(2009, 1, 2), dt.datetime(2009, 1, 5)], 'list'),
        ('D', dt.datetime(2009, 1, 2), 'date'),
        ('D', (dt.date(2009, 1, 2), dt.date(2009, 1, 5)), 'date_range'),
    ]
)
@pytest.mark.parametrize("dask", [True, False])
def test_file_source_filter(source, column_value_type, dask, expected_df):
    column, value, _ = column_value_type
    kwargs = {column: value, '__dask': dask}
    filtered = source.get('test', **kwargs)
    pd.testing.assert_frame_equal(filtered, expected_df)


@pytest.mark.parametrize(
    "column_value_type", [
        ('A', 1, 'single_value'),
        ('A', (1, 3), 'range'),
        ('A', (1, 2), 'range'),
        ('A', [(0, 1), (3, 4)], 'range_list'),
        ('C', 'foo2', 'single_value'),
        ('C', ['foo1', 'foo3'], 'list'),
        ('D', dt.datetime(2009, 1, 2), 'single_value'),
        ('D', (dt.datetime(2009, 1, 2), dt.datetime(2009, 1, 5)), 'range'),
        ('D', [dt.datetime(2009, 1, 2), dt.datetime(2009, 1, 5)], 'list'),
        ('D', dt.datetime(2009, 1, 2), 'date'),
        ('D', (dt.date(2009, 1, 2), dt.date(2009, 1, 5)), 'date_range'),
    ]
)
@pytest.mark.parametrize("dask", [True, False])
def test_file_source_get_query_cache(source, column_value_type, dask, expected_df):
    column, value, _ = column_value_type
    kwargs = {column: value}
    source.get('test', __dask=dask, **kwargs)
    cache_key = source._get_key('test', **kwargs)
    assert cache_key in source._cache
    cached_df = source._cache[cache_key]
    if dask:
        cached_df = cached_df.compute()
    pd.testing.assert_frame_equal(cached_df, expected_df)
    cache_key = source._get_key('test', **kwargs)
    assert cache_key in source._cache


@pytest.mark.parametrize(
    "column_value_type", [
        ('A', 1, 'single_value'),
        ('A', [(0, 1), (3, 4)], 'range_list'),
        ('C', ['foo1', 'foo3'], 'list'),
        ('D', (dt.datetime(2009, 1, 2), dt.datetime(2009, 1, 5)), 'range'),
        ('D', dt.datetime(2009, 1, 2), 'date'),
    ]
)
@pytest.mark.parametrize("dask", [True, False])
def test_file_source_clear_cache(source, column_value_type, dask):
    column, value, _ = column_value_type
    kwargs = {column: value}
    source.get('test', __dask=dask, **kwargs)
    cache_key = source._get_key('test', **kwargs)
    assert cache_key in source._cache
    source.clear_cache()
    assert len(source._cache) == 0


def test_file_source_get_query_cache_to_file(make_filesource, cachedir):
    root = os.path.dirname(__file__)
    source = make_filesource(root, cache_dir=cachedir)
    source.get('test', A=(1, 2))

    cache_key = source._get_key('test', A=(1, 2))
    cache_path = Path(cachedir) / f'{cache_key}_test.parq'
    df = pd.read_parquet(cache_path, engine='fastparquet')

    # Patch index names due to https://github.com/dask/fastparquet/issues/732
    df.index.names = [None]
    pd.testing.assert_frame_equal(
        df,
        makeMixedDataFrame().iloc[1:3]
    )


def test_file_source_get_tables(source):
    tables = source.get_tables()
    assert tables == ['test']


def test_file_source_variable(make_variable_filesource, mixed_df):
    root = os.path.dirname(__file__)
    source = make_variable_filesource(root)
    state.variables.tables = {'test': 'test2.csv'}
    df = source.get('test')
    expected = mixed_df.iloc[::-1].reset_index(drop=True)
    pd.testing.assert_frame_equal(df, expected)


def test_extension_of_comlicated_url(source):
    url = "https://api.tfl.gov.uk/Occupancy/BikePoints/@{stations.stations.id}?app_key=random_numbers"
    source.tables["test"] = url
    assert source._named_files["test"][1] is None
