import os
import pytest

from google.cloud.bigtable import Client
from google.cloud.bigtable import column_family
from bt_df_data_retention_manager.state.record.retention import JobStatusEnum
from bt_df_data_retention_manager.state.storage.bigtable import StateStorageManagerBigtable
from bt_df_data_retention_manager.state.storage.interface import NonexistentStateId
from bt_df_data_retention_manager.state.record.retention import TableRetentionState
from tests.test_retention_dates import example_params

column_family_id = "status"


@pytest.fixture(scope="module")
def init_local_bigtable():
    project = 'fake-project'
    instance_id = 'fake-instance'
    state_table_id = 'dataflow_data_retention_state'

    os.environ['BIGTABLE_EMULATOR_HOST'] = 'localhost:8086'

    client = Client(project=project, admin=True)
    instance = client.instance(instance_id)

    # [create dataflow_data_retention_state ]
    table = instance.table(state_table_id)

    max_versions_rule = column_family.MaxVersionsGCRule(1)

    column_families = {column_family_id: max_versions_rule}
    if not table.exists():
        table.create(column_families=column_families)
    else:
        print("Table {} already exists.".format(state_table_id))
    # [END bigtable_hw_create_table]

    return table


@pytest.fixture
def state_storage_manager(init_local_bigtable):
    return StateStorageManagerBigtable(init_local_bigtable, column_family_id)


@pytest.fixture
def sample_status_data(example_params):
    charts = TableRetentionState(table_id="charts",
                                 job_status=JobStatusEnum.ready,
                                 latest_job_id='',
                                 current_end_date=example_params.end_date,
                                 retries=0)

    chartmetric = TableRetentionState(table_id="chartmetric",
                                      job_status=JobStatusEnum.ready,
                                      latest_job_id='',
                                      current_end_date=example_params.end_date,
                                      retries=0)

    return [charts, chartmetric]


def test_state_bt_rwd(state_storage_manager, sample_status_data):
    state_storage_manager.put(sample_status_data[0])
    test_read_data = state_storage_manager.read(sample_status_data[0].table_id)
    assert test_read_data == sample_status_data[0]
    state_storage_manager.delete(sample_status_data[0].table_id)
    try:
        test_read_data = state_storage_manager.read(
            sample_status_data[0].table_id)
    except NonexistentStateId as e:
        assert e.missing_id == sample_status_data[0].table_id


def test_state_bt_create(state_storage_manager, sample_status_data):
    try:
        state_storage_manager.delete(sample_status_data[1].table_id)
    except NonexistentStateId:
        pass
    try:
        _ = state_storage_manager.read(sample_status_data[1].table_id)
    except NonexistentStateId:
        state_storage_manager.put(sample_status_data[1])

    assert sample_status_data[1] == state_storage_manager.read(
        sample_status_data[1].table_id)