# -*- coding: utf-8 -*-

import pytest

from ruamel.yaml import YAML
import os

import pandas as pd

from great_expectations.core import ExpectationSuite
from great_expectations.core.batch import Batch
from great_expectations.exceptions import BatchKwargsError
from great_expectations.datasource import SparkDFDatasource
from great_expectations.dataset import SparkDFDataset
from great_expectations.datasource.types import InMemoryBatchKwargs
from great_expectations.validator.validator import Validator

yaml = YAML()


@pytest.fixture(scope="module")
def test_folder_connection_path(tmp_path_factory):
    df1 = pd.DataFrame(
        {'col_1': [1, 2, 3, 4, 5], 'col_2': ['a', 'b', 'c', 'd', 'e']})
    path = str(tmp_path_factory.mktemp("test_folder_connection_path"))
    df1.to_csv(os.path.join(path, "test.csv"))

    return str(path)


@pytest.fixture(scope="module")
def test_parquet_folder_connection_path(tmp_path_factory):
    df1 = pd.DataFrame(
        {'col_1': [1, 2, 3, 4, 5], 'col_2': ['a', 'b', 'c', 'd', 'e']})
    basepath = str(tmp_path_factory.mktemp("parquet_context"))
    df1.to_parquet(os.path.join(basepath, "test.parquet"))

    return basepath


def test_sparkdf_datasource_custom_data_asset(data_context, test_folder_connection_path, spark_session):
    assert spark_session  # Ensure a spark session exists
    name = "test_sparkdf_datasource"
    # type_ = "spark"
    class_name = "SparkDFDatasource"

    data_asset_type_config = {
        "module_name": "custom_sparkdf_dataset",
        "class_name": "CustomSparkDFDataset"
    }
    data_context.add_datasource(name,
                                class_name=class_name,
                                data_asset_type=data_asset_type_config,
                                batch_kwargs_generators={
    "subdir_reader": {
        "class_name": "SubdirReaderBatchKwargsGenerator",
        "base_directory": test_folder_connection_path
    }
}
)


    # We should now see updated configs
    with open(os.path.join(data_context.root_directory, "great_expectations.yml"), "r") as data_context_config_file:
        data_context_file_config = yaml.load(data_context_config_file)

    assert data_context_file_config["datasources"][name]["data_asset_type"]["module_name"] == "custom_sparkdf_dataset"
    assert data_context_file_config["datasources"][name]["data_asset_type"]["class_name"] == "CustomSparkDFDataset"

    # We should be able to get a dataset of the correct type from the datasource.
    data_context.create_expectation_suite("test_sparkdf_datasource.default")
    batch_kwargs = data_context.build_batch_kwargs(name, "subdir_reader", "test")
    batch_kwargs["reader_options"] = {
            'header': True,
            'inferSchema': True
        }
    batch = data_context.get_batch(
        batch_kwargs=batch_kwargs,
        expectation_suite_name="test_sparkdf_datasource.default"
    )
    assert type(batch).__name__ == "CustomSparkDFDataset"
    res = batch.expect_column_approx_quantile_values_to_be_between("col_1", quantile_ranges={
        "quantiles": [0., 1.],
        "value_ranges": [[1, 1], [5, 5]]
    })
    assert res.success is True


def test_create_sparkdf_datasource(data_context, tmp_path_factory):
    pyspark_skip = pytest.importorskip("pyspark")
    base_dir = tmp_path_factory.mktemp('test_create_sparkdf_datasource')
    name = "test_sparkdf_datasource"
    # type_ = "spark"
    class_name = "SparkDFDatasource"

    data_context.add_datasource(name, class_name=class_name,
                                batch_kwargs_generators={
                                    "default": {
                                        "class_name": "SubdirReaderBatchKwargsGenerator",
                                        "base_directory": str(base_dir)
                                    }
                                }
                                )
    data_context_config = data_context.get_config()

    assert name in data_context_config["datasources"]
    assert data_context_config["datasources"][name]["class_name"] == class_name
    assert data_context_config["datasources"][name]["batch_kwargs_generators"]["default"]["base_directory"] == str(base_dir)

    base_dir = tmp_path_factory.mktemp('test_create_sparkdf_datasource-2')
    name = "test_sparkdf_datasource"

    data_context.add_datasource(name,
                                class_name=class_name,
                                batch_kwargs_generators={
                                    "default": {
                                        "class_name": "SubdirReaderBatchKwargsGenerator",
                                        "reader_options":
                                            {
                                                "sep": "|",
                                                "header": False
                                            }
                                    }
                                    }
                                )
    data_context_config = data_context.get_config()

    assert name in data_context_config["datasources"]
    assert data_context_config["datasources"][name]["class_name"] == class_name
    assert data_context_config["datasources"][name]["batch_kwargs_generators"]["default"]["reader_options"]["sep"] == "|"

    # Note that pipe is special in yml, so let's also check to see that it was properly serialized
    with open(os.path.join(data_context.root_directory, "great_expectations.yml"), "r") as configfile:
        lines = configfile.readlines()
        assert "          sep: '|'\n" in lines
        assert "          header: false\n" in lines


def test_standalone_spark_parquet_datasource(test_parquet_folder_connection_path, spark_session):
    assert spark_session  # Ensure a sparksession exists
    datasource = SparkDFDatasource('SparkParquet', batch_kwargs_generators={
    "subdir_reader": {
        "class_name": "SubdirReaderBatchKwargsGenerator",
        "base_directory": test_parquet_folder_connection_path
    }
}
                                   )


    assert datasource.get_available_data_asset_names()["subdir_reader"]["names"] == [('test', 'file')]
    batch = datasource.get_batch(batch_kwargs={
                                       "path": os.path.join(test_parquet_folder_connection_path,
                                                            'test.parquet')
                                   })
    assert isinstance(batch, Batch)
    # NOTE: below is a great example of CSV vs. Parquet typing: pandas reads content as string, spark as int
    assert batch.data.head()['col_1'] == 1
    assert batch.data.count() == 5

    # Limit should also work
    batch = datasource.get_batch(batch_kwargs={
                                       "path": os.path.join(test_parquet_folder_connection_path,
                                                            'test.parquet'),
                                       "limit": 2
                                   })
    assert isinstance(batch, Batch)
    # NOTE: below is a great example of CSV vs. Parquet typing: pandas reads content as string, spark as int
    assert batch.data.head()['col_1'] == 1
    assert batch.data.count() == 2


def test_standalone_spark_csv_datasource(test_folder_connection_path):
    pyspark_skip = pytest.importorskip("pyspark")
    datasource = SparkDFDatasource('SparkParquet',
                                   batch_kwargs_generators={"subdir_reader": {
                                        "class_name": "SubdirReaderBatchKwargsGenerator",
                                        "base_directory": test_folder_connection_path
                                        }
                                    }
                                   )

    assert datasource.get_available_data_asset_names()["subdir_reader"]["names"] == [('test', 'file')]
    batch = datasource.get_batch(batch_kwargs={
                                       "path": os.path.join(test_folder_connection_path,
                                                            'test.csv'),
                                       "reader_options": {"header": True}
                                   })
    assert isinstance(batch, Batch)
    # NOTE: below is a great example of CSV vs. Parquet typing: pandas reads content as string, spark as int
    assert batch.data.head()['col_1'] == '1'


def test_standalone_spark_passthrough_datasource(data_context, dataset):
    pyspark_skip = pytest.importorskip("pyspark")
    datasource = data_context.add_datasource("spark_source",
                                             module_name="great_expectations.datasource",
                                             class_name="SparkDFDatasource")

    # We want to ensure that an externally-created spark DataFrame can be successfully instantiated using the
    # datasource built in a data context
    # Our dataset fixture is parameterized by all backends. The spark source should only accept a spark dataset
    data_context.create_expectation_suite("new_suite")
    batch_kwargs = InMemoryBatchKwargs(datasource="spark_source", dataset=dataset)

    if isinstance(dataset, SparkDFDataset):
        # We should be smart enough to figure out this is a batch:
        batch = data_context.get_batch(batch_kwargs=batch_kwargs, expectation_suite_name="new_suite")
        res = batch.expect_column_to_exist("infinities")
        assert res.success is True
        res = batch.expect_column_to_exist("not_a_column")
        assert res.success is False
        batch.save_expectation_suite()
        assert os.path.isfile(os.path.join(
            data_context.root_directory,
            "expectations/new_suite.json")
        )

    else:
        with pytest.raises(BatchKwargsError) as exc:
            # noinspection PyUnusedLocal
            batch = data_context.get_batch(batch_kwargs=batch_kwargs, expectation_suite_name="new_suite")
            assert "Unrecognized batch_kwargs for spark_source" in exc.value.message


def test_invalid_reader_sparkdf_datasource(tmp_path_factory):
    pytest.importorskip("pyspark")
    basepath = str(tmp_path_factory.mktemp("test_invalid_reader_sparkdf_datasource"))
    datasource = SparkDFDatasource(
        'mysparksource',
        batch_kwargs_generators={
            "subdir_reader": {
                "class_name": "SubdirReaderBatchKwargsGenerator",
                "base_directory": basepath
            }
        }
   )

    with open(os.path.join(basepath, "idonotlooklikeacsvbutiam.notrecognized"), "w") as newfile:
        newfile.write("a,b\n1,2\n3,4\n")

    with pytest.raises(BatchKwargsError) as exc:
        datasource.get_batch(batch_kwargs={
            "path": os.path.join(basepath, "idonotlooklikeacsvbutiam.notrecognized")
        })
        assert "Unable to determine reader for path" in exc.value.message

    with pytest.raises(BatchKwargsError) as exc:
        datasource.get_batch(batch_kwargs={
            "path": os.path.join(basepath, "idonotlooklikeacsvbutiam.notrecognized"),
            "reader_method": "blarg"
        })
        assert "Unknown reader method: blarg" in exc.value.message

    with pytest.raises(BatchKwargsError) as exc:
        datasource.get_batch(batch_kwargs={
            "path": os.path.join(basepath, "idonotlooklikeacsvbutiam.notrecognized"),
            "reader_method": "excel"
        })
        assert "Unknown reader: excel" in exc.value.message

    batch = datasource.get_batch(batch_kwargs={
            "path": os.path.join(basepath, "idonotlooklikeacsvbutiam.notrecognized"),
            "reader_method": "csv",
            "reader_options": {'header': True}
    })
    assert batch.data.head()["a"] == "1"


def test_spark_config():
    pytest.importorskip("pyspark")
    source = SparkDFDatasource()
    conf = source.spark.sparkContext.getConf().getAll()
    # Without specifying any spark_config values we get defaults
    assert ("spark.app.name", "pyspark-shell") in conf

    source = SparkDFDatasource(spark_config={
        "spark.app.name": "great_expectations",
        "spark.sql.catalogImplementation": "hive",
        "spark.executor.memory": "128m"
    })

    # Test that our values were set
    conf = source.spark.sparkContext.getConf().getAll()
    assert ("spark.app.name", "great_expectations") in conf
    assert ("spark.sql.catalogImplementation", "hive") in conf
    assert ("spark.executor.memory", "128m") in conf


def test_pandas_datasource_processes_dataset_options(test_folder_connection_path):
    pytest.importorskip("pyspark")
    datasource = SparkDFDatasource('PandasCSV', batch_kwargs_generators={
            "subdir_reader": {
                "class_name": "SubdirReaderBatchKwargsGenerator",
                "base_directory": test_folder_connection_path
            }
        }
    )
    batch_kwargs = datasource.build_batch_kwargs("subdir_reader", name="test")
    batch_kwargs["dataset_options"] = {"caching": False, "persist": False}
    batch = datasource.get_batch(batch_kwargs)
    validator = Validator(batch, ExpectationSuite(expectation_suite_name="foo"))
    dataset = validator.get_dataset()
    assert dataset.caching is False
    assert dataset._persist is False
