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

# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import importlib

from google.api_core import exceptions
from google.api_core import operation
from google.cloud import bigquery
from google.cloud import bigquery_storage
from google.cloud import aiplatform
from google.cloud.aiplatform.vertex_ray import bigquery_datasource
from google.cloud.aiplatform.vertex_ray.bigquery_datasink import (
    _BigQueryDatasink,
)
import test_constants as tc
from google.cloud.bigquery import job
from google.cloud.bigquery_storage_v1.types import stream as gcbqs_stream
import mock
import pytest
import pyarrow as pa
import ray


_TEST_BQ_DATASET_ID = "mockdataset"
_TEST_BQ_TABLE_ID = "mocktable"
_TEST_BQ_DATASET = _TEST_BQ_DATASET_ID + "." + _TEST_BQ_TABLE_ID
_TEST_BQ_TEMP_DESTINATION = (
    tc.ProjectConstants.TEST_GCP_PROJECT_ID + ".tempdataset.temptable"
)
_TEST_DISPLAY_NAME = "display_name"


@pytest.fixture(autouse=True)
def bq_client_full_mock(monkeypatch):
    client_mock = mock.create_autospec(bigquery.Client)
    client_mock.return_value = client_mock

    def bq_get_dataset_mock(dataset_id):
        if dataset_id != _TEST_BQ_DATASET_ID:
            raise exceptions.NotFound(
                "[Ray on Vertex AI]: Dataset {} is not found. Please ensure that it"
                " exists.".format(_TEST_BQ_DATASET)
            )

    def bq_get_table_mock(table_id):
        if table_id != _TEST_BQ_DATASET:
            raise exceptions.NotFound(
                "[Ray on Vertex AI]: Table {} is not found. Please ensure that it"
                " exists.".format(_TEST_BQ_DATASET)
            )

    def bq_create_dataset_mock(dataset_id, **kwargs):
        if dataset_id == "existingdataset":
            raise exceptions.Conflict("Dataset already exists")
        return mock.Mock(operation.Operation)

    def bq_delete_table_mock(table, **kwargs):
        return None

    def bq_query_mock(query):
        fake_job_ref = job._JobReference(
            "fake_job_id",
            tc.ProjectConstants.TEST_GCP_PROJECT_ID,
            "us-central1",
        )
        fake_query_job = job.QueryJob(fake_job_ref, query, None)
        try:
            fake_query_job.configuration.destination = _TEST_BQ_TEMP_DESTINATION
        except AttributeError:
            fake_query_job._configuration.destination = _TEST_BQ_TEMP_DESTINATION
        return fake_query_job

    client_mock.get_dataset = bq_get_dataset_mock
    client_mock.get_table = bq_get_table_mock
    client_mock.create_dataset = bq_create_dataset_mock
    client_mock.delete_table = bq_delete_table_mock
    client_mock.query = bq_query_mock

    monkeypatch.setattr(bigquery, "Client", client_mock)
    client_mock.reset_mock()
    return client_mock


@pytest.fixture(autouse=True)
def bqs_client_full_mock(monkeypatch):
    client_mock = mock.create_autospec(bigquery_storage.BigQueryReadClient)
    client_mock.return_value = client_mock

    def bqs_create_read_session(max_stream_count=0, **kwargs):
        read_session_proto = gcbqs_stream.ReadSession()
        read_session_proto.streams = [
            gcbqs_stream.ReadStream() for _ in range(max_stream_count)
        ]
        return read_session_proto

    client_mock.create_read_session = bqs_create_read_session

    monkeypatch.setattr(bigquery_storage, "BigQueryReadClient", client_mock)
    client_mock.reset_mock()
    return client_mock


@pytest.fixture
def bq_query_result_mock():
    with mock.patch.object(bigquery.job.QueryJob, "result") as query_result_mock:
        yield query_result_mock


@pytest.fixture
def bq_query_result_mock_fail():
    with mock.patch.object(bigquery.job.QueryJob, "result") as query_result_mock_fail:
        query_result_mock_fail.side_effect = exceptions.BadRequest("400 Syntax error")
        yield query_result_mock_fail


@pytest.fixture
def ray_remote_function_mock():
    with mock.patch.object(ray.remote_function.RemoteFunction, "_remote") as remote_fn:
        remote_fn.return_value = 1
        yield remote_fn


@pytest.fixture
def ray_get_mock():
    with mock.patch.object(ray, "get") as ray_get:
        ray_get.return_value = None
        yield ray_get


class TestReadBigQuery:
    """Tests for BigQuery Read."""

    def setup_method(self):
        importlib.reload(aiplatform.initializer)
        importlib.reload(aiplatform)

    def teardown_method(self):
        aiplatform.initializer.global_pool.shutdown(wait=True)

    @pytest.mark.parametrize(
        "parallelism",
        [1, 2, 3, 4, 10, 100],
    )
    def test_create_reader(self, parallelism):
        bq_ds = bigquery_datasource._BigQueryDatasource(
            project_id=tc.ProjectConstants.TEST_GCP_PROJECT_ID,
            dataset=_TEST_BQ_DATASET,
        )
        read_tasks_list = bq_ds.get_read_tasks(parallelism)
        assert len(read_tasks_list) == parallelism

    @pytest.mark.parametrize(
        "parallelism",
        [1, 2, 3, 4, 10, 100],
    )
    def test_create_reader_initialized(self, parallelism):
        """If initialized, create_reader doesn't need to specify project_id."""
        aiplatform.init(
            project=tc.ProjectConstants.TEST_GCP_PROJECT_ID,
            staging_bucket=tc.ProjectConstants.TEST_ARTIFACT_URI,
        )
        bq_ds = bigquery_datasource._BigQueryDatasource(
            dataset=_TEST_BQ_DATASET,
        )
        read_tasks_list = bq_ds.get_read_tasks(parallelism)
        assert len(read_tasks_list) == parallelism

    @pytest.mark.parametrize(
        "parallelism",
        [1, 2, 3, 4, 10, 100],
    )
    def test_create_reader_query(self, parallelism, bq_query_result_mock):
        bq_ds = bigquery_datasource._BigQueryDatasource(
            project_id=tc.ProjectConstants.TEST_GCP_PROJECT_ID,
            query="SELECT * FROM mockdataset.mocktable",
        )
        read_tasks_list = bq_ds.get_read_tasks(parallelism)
        bq_query_result_mock.assert_called_once()
        assert len(read_tasks_list) == parallelism

    @pytest.mark.parametrize(
        "parallelism",
        [1, 2, 3, 4, 10, 100],
    )
    def test_create_reader_query_bad_request(
        self,
        parallelism,
        bq_query_result_mock_fail,
    ):
        bq_ds = bigquery_datasource._BigQueryDatasource(
            project_id=tc.ProjectConstants.TEST_GCP_PROJECT_ID,
            query="SELECT * FROM mockdataset.mocktable",
        )
        with pytest.raises(exceptions.BadRequest):
            bq_ds.get_read_tasks(parallelism)
        bq_query_result_mock_fail.assert_called()

    def test_dataset_query_kwargs_provided(self):
        with pytest.raises(ValueError) as exception:
            bigquery_datasource._BigQueryDatasource(
                project_id=tc.ProjectConstants.TEST_GCP_PROJECT_ID,
                dataset=_TEST_BQ_DATASET,
                query="SELECT * FROM mockdataset.mocktable",
            )
        expected_message = (
            "[Ray on Vertex AI]: Query and dataset kwargs cannot both be provided"
            " (must be mutually exclusive)."
        )
        assert str(exception.value) == expected_message

    def test_create_reader_dataset_not_found(self):
        parallelism = 4
        bq_ds = bigquery_datasource._BigQueryDatasource(
            project_id=tc.ProjectConstants.TEST_GCP_PROJECT_ID,
            dataset="nonexistentdataset.mocktable",
        )
        with pytest.raises(ValueError) as exception:
            bq_ds.get_read_tasks(parallelism)
        expected_message = (
            "[Ray on Vertex AI]: Dataset nonexistentdataset is not found. Please"
            " ensure that it exists."
        )
        assert str(exception.value) == expected_message

    def test_create_reader_table_not_found(self):
        parallelism = 4
        bq_ds = bigquery_datasource._BigQueryDatasource(
            project_id=tc.ProjectConstants.TEST_GCP_PROJECT_ID,
            dataset="mockdataset.nonexistenttable",
        )
        with pytest.raises(ValueError) as exception:
            bq_ds.get_read_tasks(parallelism)
        expected_message = (
            "[Ray on Vertex AI]: Table mockdataset.nonexistenttable is not found."
            " Please ensure that it exists."
        )
        assert str(exception.value) == expected_message


@pytest.mark.usefixtures("google_auth_mock")
class TestWriteBigQuery:
    """Tests for BigQuery Write."""

    def setup_method(self):
        importlib.reload(aiplatform.initializer)
        importlib.reload(aiplatform)

    def teardown_method(self):
        aiplatform.initializer.global_pool.shutdown(wait=True)

    def test_write(self, ray_get_mock, ray_remote_function_mock):
        if _BigQueryDatasink is None:
            return
        bq_datasink = _BigQueryDatasink(
            project_id=tc.ProjectConstants.TEST_GCP_PROJECT_ID,
            dataset=_TEST_BQ_DATASET,
        )
        arr = pa.array([2, 4, 5, 100])
        block = pa.Table.from_arrays([arr], names=["data"])
        status = bq_datasink.write(
            blocks=[block],
            ctx=None,
        )
        assert status == "ok"

    def test_write_dataset_exists(self, ray_get_mock, ray_remote_function_mock):
        if _BigQueryDatasink is None:
            return
        bq_datasink = _BigQueryDatasink(
            project_id=tc.ProjectConstants.TEST_GCP_PROJECT_ID,
            dataset="existingdataset" + "." + _TEST_BQ_TABLE_ID,
        )
        arr = pa.array([2, 4, 5, 100])
        block = pa.Table.from_arrays([arr], names=["data"])
        status = bq_datasink.write(
            blocks=[block],
            ctx=None,
        )
        assert status == "ok"
