# -*- coding: utf-8 -*-
#
# Copyright 2019 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
#
#     https://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.
"""System tests for reading rows from tables."""

import copy
import datetime as dt
import decimal
import re

import pytest
import pytz

from google.cloud import bigquery


def _to_bq_table_ref(table_name_string, partition_suffix=""):
    """Converts protobuf table reference to bigquery table reference.

    Args:
        table_name_string (str):
            A table name in storage API format.
            `projects/<projectid>/datasets/<datasetid>/tables/<tableid>`
        partition_suffix (str):
            An optional suffix to append to the table_id, useful for selecting
            partitions of ingestion-time partitioned tables.

    Returns:
        google.cloud.bigquery.table.TableReference
    """
    parts = table_name_string.split("/")

    return bigquery.table.TableReference.from_api_repr(
        {
            "projectId": parts[1],
            "datasetId": parts[3],
            "tableId": parts[5] + partition_suffix,
        }
    )


@pytest.mark.parametrize(
    "data_format,expected_schema_type",
    (("AVRO", "avro_schema"), ("ARROW", "arrow_schema")),
)
def test_read_rows_as_blocks_full_table(
    client_and_types,
    project_id,
    small_table_reference,
    data_format,
    expected_schema_type,
):
    client, types = client_and_types
    read_session = types.ReadSession()
    read_session.table = small_table_reference
    read_session.data_format = data_format

    session = client.create_read_session(
        request={
            "parent": "projects/{}".format(project_id),
            "read_session": read_session,
            "max_stream_count": 1,
        }
    )
    stream = session.streams[0].name

    schema_type = session._pb.WhichOneof("schema")
    assert schema_type == expected_schema_type

    blocks = list(client.read_rows(stream))

    assert len(blocks) > 0


@pytest.mark.parametrize("data_format", ("AVRO", "ARROW"))
def test_read_rows_as_rows_full_table(
    client_and_types, project_id, small_table_reference, data_format
):
    client, types = client_and_types
    read_session = types.ReadSession()
    read_session.table = small_table_reference
    read_session.data_format = data_format

    session = client.create_read_session(
        request={
            "parent": "projects/{}".format(project_id),
            "read_session": read_session,
            "max_stream_count": 1,
        }
    )
    stream = session.streams[0].name

    rows = list(client.read_rows(stream).rows(session))

    assert len(rows) > 0


@pytest.mark.parametrize("data_format", ("AVRO", "ARROW"))
def test_basic_nonfiltered_read(
    client_and_types, project_id, table_with_data_ref, data_format
):
    client, types = client_and_types
    read_session = types.ReadSession()
    read_session.table = table_with_data_ref
    read_session.data_format = data_format

    session = client.create_read_session(
        request={
            "parent": "projects/{}".format(project_id),
            "read_session": read_session,
            "max_stream_count": 1,
        }
    )
    stream = session.streams[0].name

    rows = list(client.read_rows(stream).rows(session))

    assert len(rows) == 5  # all table rows


def test_filtered_rows_read(client_and_types, project_id, table_with_data_ref):
    client, types = client_and_types
    read_session = types.ReadSession()
    read_session.table = table_with_data_ref
    read_session.data_format = types.DataFormat.AVRO
    read_session.read_options.row_restriction = "age >= 50"

    session = client.create_read_session(
        request={
            "parent": "projects/{}".format(project_id),
            "read_session": read_session,
            "max_stream_count": 1,
        }
    )
    stream = session.streams[0].name

    rows = list(client.read_rows(stream).rows(session))

    assert len(rows) == 2


@pytest.mark.parametrize("data_format", ("AVRO", "ARROW"))
def test_column_selection_read(
    client_and_types, project_id, table_with_data_ref, data_format
):
    client, types = client_and_types
    read_session = types.ReadSession()
    read_session.table = table_with_data_ref
    read_session.data_format = data_format
    read_session.read_options.selected_fields.append("first_name")
    read_session.read_options.selected_fields.append("age")

    session = client.create_read_session(
        request={
            "parent": "projects/{}".format(project_id),
            "read_session": read_session,
            "max_stream_count": 1,
        }
    )
    stream = session.streams[0].name

    rows = list(client.read_rows(stream).rows(session))

    for row in rows:
        assert sorted(row.keys()) == ["age", "first_name"]


def test_snapshot(client_and_types, project_id, table_with_data_ref, bq_client):
    client, types = client_and_types
    before_new_data = dt.datetime.now(tz=dt.timezone.utc)

    # load additional data into the table
    new_data = [
        {"first_name": "NewGuyFoo", "last_name": "Smith", "age": 46},
        {"first_name": "NewGuyBar", "last_name": "Jones", "age": 30},
    ]

    destination = _to_bq_table_ref(table_with_data_ref)
    bq_client.load_table_from_json(new_data, destination).result()

    # read data using the timestamp before the additional data load

    read_session = types.ReadSession()
    read_session.table = table_with_data_ref
    read_session.table_modifiers.snapshot_time = before_new_data
    read_session.data_format = types.DataFormat.AVRO

    session = client.create_read_session(
        request={
            "parent": "projects/{}".format(project_id),
            "read_session": read_session,
            "max_stream_count": 1,
        }
    )
    stream = session.streams[0].name

    rows = list(client.read_rows(stream).rows(session))

    # verify that only the data before the timestamp was returned
    assert len(rows) == 5  # all initial records

    for row in rows:
        assert "NewGuy" not in row["first_name"]  # no new records


def test_column_partitioned_table(
    client_and_types, project_id, col_partition_table_ref, bq_client
):
    client, types = client_and_types
    data = [
        {"description": "Tracking established.", "occurred": "2017-02-15"},
        {"description": "Look, a solar eclipse!", "occurred": "2018-02-15"},
        {"description": "Fake solar eclipse reported.", "occurred": "2018-02-15"},
        {"description": "1 day after false eclipse report.", "occurred": "2018-02-16"},
        {"description": "1 year after false eclipse report.", "occurred": "2019-02-15"},
    ]

    destination = _to_bq_table_ref(col_partition_table_ref)
    bq_client.load_table_from_json(data, destination).result()

    # Read from the table with a partition filter specified, and verify that
    # only the expected data is returned.

    read_session = types.ReadSession()
    read_session.table = col_partition_table_ref
    read_session.data_format = types.DataFormat.AVRO
    read_session.read_options.row_restriction = "occurred = '2018-02-15'"

    session = client.create_read_session(
        request={
            "parent": "projects/{}".format(project_id),
            "read_session": read_session,
            "max_stream_count": 1,
        }
    )

    assert session.streams  # there should be some data to fetch

    stream = session.streams[0].name

    rows = list(client.read_rows(stream).rows(session))
    assert len(rows) == 2

    expected_descriptions = ("Look, a solar eclipse!", "Fake solar eclipse reported.")
    for row in rows:
        assert row["occurred"] == dt.date(2018, 2, 15)
        assert row["description"] in expected_descriptions


@pytest.mark.parametrize("data_format", ("AVRO", "ARROW"))
def test_ingestion_time_partitioned_table(
    client_and_types, project_id, ingest_partition_table_ref, bq_client, data_format
):
    client, types = client_and_types
    data = [{"shape": "cigar", "altitude": 1200}, {"shape": "disc", "altitude": 750}]
    destination = _to_bq_table_ref(
        ingest_partition_table_ref, partition_suffix="$20190809"
    )
    bq_client.load_table_from_json(data, destination).result()

    data = [
        {"shape": "sphere", "altitude": 3500},
        {"shape": "doughnut", "altitude": 100},
    ]
    destination = _to_bq_table_ref(
        ingest_partition_table_ref, partition_suffix="$20190810"
    )
    bq_client.load_table_from_json(data, destination).result()

    data = [
        {"shape": "elephant", "altitude": 1},
        {"shape": "rocket", "altitude": 12700},
    ]
    destination = _to_bq_table_ref(
        ingest_partition_table_ref, partition_suffix="$20190811"
    )
    bq_client.load_table_from_json(data, destination).result()

    read_session = types.ReadSession()
    read_session.table = ingest_partition_table_ref
    read_session.data_format = data_format
    read_session.read_options.row_restriction = "DATE(_PARTITIONTIME) = '2019-08-10'"

    session = client.create_read_session(
        request={
            "parent": "projects/{}".format(project_id),
            "read_session": read_session,
            "max_stream_count": 1,
        }
    )

    assert session.streams  # there should be some data to fetch

    stream = session.streams[0].name

    rows = list(client.read_rows(stream).rows(session))
    assert len(rows) == 2

    data_format = getattr(types.DataFormat, data_format)
    if data_format == types.DataFormat.AVRO:
        actual_items = {(row["shape"], row["altitude"]) for row in rows}
    elif data_format == types.DataFormat.ARROW:
        actual_items = {(row["shape"].as_py(), row["altitude"].as_py()) for row in rows}
    else:
        raise AssertionError(f"got unexpected data_format: {data_format}")

    expected_items = {("sphere", 3500), ("doughnut", 100)}
    assert actual_items == expected_items


@pytest.mark.parametrize("data_format", ("AVRO", "ARROW"))
def test_decoding_data_types(
    client_and_types, project_id, all_types_table_ref, bq_client, data_format
):
    client, types = client_and_types
    data = [
        {
            "string_field": "Price: € 9.95.",
            "bytes_field": bigquery._helpers._bytes_to_json(b"byteees"),
            "int64_field": -1085,
            "float64_field": -42.195,
            "numeric_field": "1.4142",
            "bool_field": True,
            "geography_field": '{"type": "Point", "coordinates": [-49.3028, 69.0622]}',
            "person_struct_field": {"name": "John", "age": 42},
            "timestamp_field": 1565357902.017896,  # 2019-08-09T13:38:22.017896
            "date_field": "1995-03-17",
            "time_field": "16:24:51",
            "datetime_field": "2005-10-26T19:49:41",
            "string_array_field": ["foo", "bar", "baz"],
        }
    ]

    # Explicit schema is needed to recognize bytes_field as BYTES, and not STRING.
    # Since partial schemas are not supported in load_table_from_json(), a full
    # schema needs to be specified.
    schema = [
        bigquery.SchemaField("string_field", "STRING"),
        bigquery.SchemaField("bytes_field", "BYTES"),
        bigquery.SchemaField("int64_field", "INT64"),
        bigquery.SchemaField("float64_field", "FLOAT64"),
        bigquery.SchemaField("numeric_field", "NUMERIC"),
        bigquery.SchemaField("bool_field", "BOOL"),
        bigquery.SchemaField("geography_field", "GEOGRAPHY"),
        bigquery.SchemaField(
            "person_struct_field",
            "STRUCT",
            fields=(
                bigquery.SchemaField("name", "STRING"),
                bigquery.SchemaField("age", "INT64"),
            ),
        ),
        bigquery.SchemaField("timestamp_field", "TIMESTAMP"),
        bigquery.SchemaField("date_field", "DATE"),
        bigquery.SchemaField("time_field", "TIME"),
        bigquery.SchemaField("datetime_field", "DATETIME"),
        bigquery.SchemaField("string_array_field", "STRING", mode="REPEATED"),
    ]

    job_config = bigquery.LoadJobConfig(schema=schema)
    destination = _to_bq_table_ref(all_types_table_ref)
    bq_client.load_table_from_json(data, destination, job_config=job_config).result()

    read_session = types.ReadSession()
    read_session.table = all_types_table_ref
    read_session.data_format = data_format

    session = client.create_read_session(
        request={
            "parent": "projects/{}".format(project_id),
            "read_session": read_session,
            "max_stream_count": 1,
        }
    )

    assert session.streams  # there should be data available

    stream = session.streams[0].name

    data_format = getattr(types.DataFormat, data_format)
    if data_format == types.DataFormat.AVRO:
        rows = list(client.read_rows(stream).rows(session))
    elif data_format == types.DataFormat.ARROW:
        rows = list(
            dict((key, value.as_py()) for key, value in row_dict.items())
            for row_dict in client.read_rows(stream).rows(session)
        )
    else:
        raise AssertionError(f"got unexpected data_format: {data_format}")

    expected_result = {
        "string_field": "Price: € 9.95.",
        "bytes_field": b"byteees",
        "int64_field": -1085,
        "float64_field": -42.195,
        "numeric_field": decimal.Decimal("1.4142"),
        "bool_field": True,
        "geography_field": "POINT(-49.3028 69.0622)",
        "person_struct_field": {"name": "John", "age": 42},
        "timestamp_field": dt.datetime(2019, 8, 9, 13, 38, 22, 17896, tzinfo=pytz.UTC),
        "date_field": dt.date(1995, 3, 17),
        "time_field": dt.time(16, 24, 51),
        "string_array_field": ["foo", "bar", "baz"],
    }

    result_copy = copy.copy(rows[0])
    del result_copy["datetime_field"]
    assert result_copy == expected_result

    # Compare datetime separately, AVRO and PYARROW return different object types,
    # although they should both represent the same value.
    # TODO: when fixed, change assertion to assert a datetime instance!
    expected_pattern = re.compile(r"2005-10-26( |T)19:49:41")
    assert expected_pattern.match(str(rows[0]["datetime_field"]))


@pytest.mark.parametrize("data_format", ("AVRO", "ARROW"))
def test_resuming_read_from_offset(
    client_and_types, project_id, data_format, local_shakespeare_table_reference
):
    client, types = client_and_types
    read_session = types.ReadSession()
    read_session.table = local_shakespeare_table_reference
    read_session.data_format = data_format

    session = client.create_read_session(
        request={
            "parent": "projects/{}".format(project_id),
            "read_session": read_session,
            "max_stream_count": 1,
        }
    )

    assert session.streams  # there should be data available

    stream = session.streams[0].name

    read_rows_stream = client.read_rows(stream)

    # fetch the first two batches of rows
    rows_iter = iter(read_rows_stream)
    some_rows = next(rows_iter)
    more_rows = next(rows_iter)

    # fetch the rest of the rows using the stream offset
    offset = some_rows.row_count + more_rows.row_count
    remaining_rows_count = sum(
        1 for _ in client.read_rows(stream, offset=offset).rows(session)
    )

    # verify that the counts match
    expected_len = 164656  # total rows in shakespeare table
    actual_len = remaining_rows_count + some_rows.row_count + more_rows.row_count
    assert actual_len == expected_len


def test_read_rows_to_dataframe_with_wide_table(client_and_types, project_id):
    # Use a wide table to boost the chance of getting a large message size.
    # https://github.com/googleapis/python-bigquery-storage/issues/78
    client, types = client_and_types
    read_session = types.ReadSession()
    read_session.table = "projects/{}/datasets/{}/tables/{}".format(
        "bigquery-public-data", "geo_census_tracts", "us_census_tracts_national"
    )
    read_session.data_format = types.DataFormat.ARROW

    session = client.create_read_session(
        request={
            "parent": "projects/{}".format(project_id),
            "read_session": read_session,
            "max_stream_count": 1,
        }
    )

    stream = session.streams[0].name

    read_rows_stream = client.read_rows(stream)

    # fetch the first two batches of rows
    pages_iter = iter(read_rows_stream.rows(session).pages)
    some_rows = next(pages_iter)

    assert all(len(row["tract_geom"].as_py()) > 0 for row in some_rows)
