"""Tests for dataset.compute_signal()."""

from datetime import datetime
from typing import Iterable

import pytest
from freezegun import freeze_time
from pandas import Timestamp
from pytest_mock import MockerFixture

from ..schema import PATH_WILDCARD, ROWID, Item, field, schema
from ..source import clear_source_registry, register_source
from .dataset import DELETED_LABEL_NAME, DatasetManifest, SelectGroupsResult, SortOrder
from .dataset_test_utils import TestDataMaker, TestSource

TEST_ITEMS: list[Item] = [{'str': 'a', 'int': 1}, {'str': 'b', 'int': 2}, {'str': 'c', 'int': 3}]

TEST_TIME = datetime(2023, 8, 15, 1, 23, 45)


@pytest.fixture(scope='module', autouse=True)
def setup_teardown() -> Iterable[None]:
  # Setup.
  register_source(TestSource)

  # Unit test runs.
  yield

  # Teardown.
  clear_source_registry()


@freeze_time(TEST_TIME)
def test_add_single_label(make_test_data: TestDataMaker, mocker: MockerFixture) -> None:
  dataset = make_test_data(TEST_ITEMS)

  num_labels = dataset.add_labels('test_label', filters=[(ROWID, 'equals', '1')])
  assert num_labels == 1
  assert dataset.manifest() == DatasetManifest(
    source=TestSource(),
    namespace='test_namespace',
    dataset_name='test_dataset',
    data_schema=schema(
      {
        'str': 'string',
        'int': 'int32',
        'test_label': field(fields={'label': 'string', 'created': 'timestamp'}, label='test_label'),
      }
    ),
    num_items=3,
  )

  assert list(dataset.select_rows([PATH_WILDCARD])) == [
    {
      'str': 'a',
      'int': 1,
      'test_label.label': 'true',
      'test_label.created': Timestamp(TEST_TIME),
    },
    {
      'str': 'b',
      'int': 2,
      'test_label.label': None,
      'test_label.created': None,
    },
    {
      'str': 'c',
      'int': 3,
      'test_label.label': None,
      'test_label.created': None,
    },
  ]

  # Make sure we can filter by the label.
  assert list(
    dataset.select_rows([PATH_WILDCARD], filters=[('test_label.label', 'equals', 'true')])
  ) == [
    {
      'str': 'a',
      'int': 1,
      'test_label.label': 'true',
      'test_label.created': Timestamp(TEST_TIME),
    },
  ]

  assert dataset.get_label_names() == ['test_label']


@freeze_time(TEST_TIME)
def test_add_labels_sort_limit(make_test_data: TestDataMaker, mocker: MockerFixture) -> None:
  dataset = make_test_data(TEST_ITEMS)

  num_labels = dataset.add_labels('test_label', sort_by=['int'], sort_order=SortOrder.ASC, limit=2)
  assert num_labels == 2
  assert dataset.manifest() == DatasetManifest(
    source=TestSource(),
    namespace='test_namespace',
    dataset_name='test_dataset',
    data_schema=schema(
      {
        'str': 'string',
        'int': 'int32',
        'test_label': field(fields={'label': 'string', 'created': 'timestamp'}, label='test_label'),
      }
    ),
    num_items=3,
  )

  assert list(dataset.select_rows([PATH_WILDCARD])) == [
    {
      'str': 'a',
      'int': 1,
      'test_label.label': 'true',
      'test_label.created': Timestamp(TEST_TIME),
    },
    {
      'str': 'b',
      'int': 2,
      'test_label.label': 'true',
      'test_label.created': Timestamp(TEST_TIME),
    },
    {
      'str': 'c',
      'int': 3,
      'test_label.label': None,
      'test_label.created': None,
    },
  ]

  assert dataset.get_label_names() == ['test_label']


@freeze_time(TEST_TIME)
def test_add_row_labels(make_test_data: TestDataMaker, mocker: MockerFixture) -> None:
  dataset = make_test_data(TEST_ITEMS)

  num_labels = dataset.add_labels('test_label', row_ids=['1', '2'])
  assert num_labels == 2
  assert list(dataset.select_rows([PATH_WILDCARD])) == [
    {
      'str': 'a',
      'int': 1,
      'test_label.label': 'true',
      'test_label.created': Timestamp(TEST_TIME),
    },
    {
      'str': 'b',
      'int': 2,
      'test_label.label': 'true',
      'test_label.created': Timestamp(TEST_TIME),
    },
    {
      'str': 'c',
      'int': 3,
      'test_label.label': None,
      'test_label.created': None,
    },
  ]
  assert dataset.get_label_names() == ['test_label']


@freeze_time(TEST_TIME)
def test_label_deleted_interaction(make_test_data: TestDataMaker, mocker: MockerFixture) -> None:
  dataset = make_test_data(TEST_ITEMS)

  # Expected semantics:
  # Adding/removing labels on deleted rows will fail unless include_deleted=True.
  # That's true even if the row was manually specified in row_ids!
  dataset.add_labels('test_label', row_ids=['1', '2'])
  dataset.delete_rows(['1'])
  dataset.remove_labels('test_label', row_ids=['1', '2'])
  assert list(dataset.select_rows([PATH_WILDCARD], include_deleted=True)) == [
    {
      'str': 'a',
      'int': 1,
      f'{DELETED_LABEL_NAME}.label': 'true',
      f'{DELETED_LABEL_NAME}.created': Timestamp(TEST_TIME),
      'test_label.label': 'true',
      'test_label.created': Timestamp(TEST_TIME),
    },
    {
      'str': 'b',
      'int': 2,
      f'{DELETED_LABEL_NAME}.label': None,
      f'{DELETED_LABEL_NAME}.created': None,
      'test_label.label': None,
      'test_label.created': None,
    },
    {
      'str': 'c',
      'int': 3,
      f'{DELETED_LABEL_NAME}.label': None,
      f'{DELETED_LABEL_NAME}.created': None,
      'test_label.label': None,
      'test_label.created': None,
    },
  ]

  dataset.remove_labels('test_label', row_ids=['1'], include_deleted=True)
  assert sorted(
    dataset.select_rows([PATH_WILDCARD], include_deleted=True), key=lambda x: x['int']
  ) == [
    {
      'str': 'a',
      'int': 1,
      f'{DELETED_LABEL_NAME}.label': 'true',
      f'{DELETED_LABEL_NAME}.created': Timestamp(TEST_TIME),
    },
    {
      'str': 'b',
      'int': 2,
      f'{DELETED_LABEL_NAME}.label': None,
      f'{DELETED_LABEL_NAME}.created': None,
    },
    {
      'str': 'c',
      'int': 3,
      f'{DELETED_LABEL_NAME}.label': None,
      f'{DELETED_LABEL_NAME}.created': None,
    },
  ]

  dataset.add_labels('test_label', row_ids=['1', '3'])
  assert sorted(
    dataset.select_rows([PATH_WILDCARD], include_deleted=True), key=lambda x: x['int']
  ) == [
    {
      'str': 'a',
      'int': 1,
      f'{DELETED_LABEL_NAME}.label': 'true',
      f'{DELETED_LABEL_NAME}.created': Timestamp(TEST_TIME),
      'test_label.label': None,
      'test_label.created': None,
    },
    {
      'str': 'b',
      'int': 2,
      f'{DELETED_LABEL_NAME}.label': None,
      f'{DELETED_LABEL_NAME}.created': None,
      'test_label.label': None,
      'test_label.created': None,
    },
    {
      'str': 'c',
      'int': 3,
      f'{DELETED_LABEL_NAME}.label': None,
      f'{DELETED_LABEL_NAME}.created': None,
      'test_label.label': 'true',
      'test_label.created': Timestamp(TEST_TIME),
    },
  ]


@freeze_time(TEST_TIME)
def test_add_row_labels_no_filters(make_test_data: TestDataMaker, mocker: MockerFixture) -> None:
  dataset = make_test_data(TEST_ITEMS)

  num_labels = dataset.add_labels('test_label')
  assert num_labels == 3
  assert list(dataset.select_rows([PATH_WILDCARD])) == [
    {
      'str': 'a',
      'int': 1,
      'test_label.label': 'true',
      'test_label.created': TEST_TIME,
    },
    {
      'str': 'b',
      'int': 2,
      'test_label.label': 'true',
      'test_label.created': TEST_TIME,
    },
    {
      'str': 'c',
      'int': 3,
      'test_label.label': 'true',
      'test_label.created': TEST_TIME,
    },
  ]

  assert dataset.get_label_names() == ['test_label']


@freeze_time(TEST_TIME)
def test_remove_labels(make_test_data: TestDataMaker, mocker: MockerFixture) -> None:
  dataset = make_test_data(TEST_ITEMS)

  num_labels = dataset.add_labels('test_label', row_ids=['1', '2', '3'])
  assert num_labels == 3

  # Remove the first label.
  num_labels = dataset.remove_labels('test_label', row_ids=['1'])
  assert num_labels == 1

  assert list(dataset.select_rows([PATH_WILDCARD], sort_by=['int'], sort_order=SortOrder.ASC)) == [
    {
      'str': 'a',
      'int': 1,
      'test_label.label': None,
      'test_label.created': None,
    },
    {
      'str': 'b',
      'int': 2,
      'test_label.label': 'true',
      'test_label.created': Timestamp(TEST_TIME),
    },
    {
      'str': 'c',
      'int': 3,
      'test_label.label': 'true',
      'test_label.created': Timestamp(TEST_TIME),
    },
  ]

  # Remove by a filter.
  dataset.remove_labels('test_label', filters=[('int', 'greater_equal', 2)])
  assert list(dataset.select_rows([PATH_WILDCARD], sort_by=['int'], sort_order=SortOrder.ASC)) == [
    {'str': 'a', 'int': 1},
    {'str': 'b', 'int': 2},
    {'str': 'c', 'int': 3},
  ]

  assert dataset.get_label_names() == []


@freeze_time(TEST_TIME)
def test_remove_labels_no_filters(make_test_data: TestDataMaker, mocker: MockerFixture) -> None:
  dataset = make_test_data(TEST_ITEMS)

  # Add all labels.
  num_labels = dataset.add_labels('test_label')
  assert num_labels == 3

  # Remove all labels.
  num_labels = dataset.remove_labels('test_label')
  assert num_labels == 3

  assert list(dataset.select_rows([PATH_WILDCARD], sort_by=['int'], sort_order=SortOrder.ASC)) == [
    {'str': 'a', 'int': 1},
    {'str': 'b', 'int': 2},
    {'str': 'c', 'int': 3},
  ]

  assert dataset.get_label_names() == []


@freeze_time(TEST_TIME)
def test_remove_labels_sort_limit(make_test_data: TestDataMaker, mocker: MockerFixture) -> None:
  dataset = make_test_data(TEST_ITEMS)

  # Add labels to every row.
  num_labels = dataset.add_labels('test_label')
  assert num_labels == 3

  # Remove 2 labels.
  num_labels = dataset.remove_labels(
    'test_label', sort_by=['int'], sort_order=SortOrder.ASC, limit=2
  )
  assert num_labels == 2
  assert dataset.manifest() == DatasetManifest(
    source=TestSource(),
    namespace='test_namespace',
    dataset_name='test_dataset',
    data_schema=schema(
      {
        'str': 'string',
        'int': 'int32',
        'test_label': field(fields={'label': 'string', 'created': 'timestamp'}, label='test_label'),
      }
    ),
    num_items=3,
  )

  assert list(dataset.select_rows([PATH_WILDCARD], sort_by=('str',), sort_order=SortOrder.ASC)) == [
    {
      'str': 'a',
      'int': 1,
      'test_label.label': None,
      'test_label.created': None,
    },
    {
      'str': 'b',
      'int': 2,
      'test_label.label': None,
      'test_label.created': None,
    },
    {
      'str': 'c',
      'int': 3,
      'test_label.label': 'true',
      'test_label.created': Timestamp(TEST_TIME),
    },
  ]

  assert dataset.get_label_names() == ['test_label']


@freeze_time(TEST_TIME)
def test_label_overwrites(make_test_data: TestDataMaker, mocker: MockerFixture) -> None:
  dataset = make_test_data(TEST_ITEMS)

  num_labels = dataset.add_labels('test_label', value='yes', filters=[(ROWID, 'equals', '1')])
  assert num_labels == 1
  # Overwrite the value.
  num_labels = dataset.add_labels('test_label', value='no', filters=[(ROWID, 'equals', '1')])
  assert num_labels == 1

  assert list(dataset.select_rows([PATH_WILDCARD])) == [
    {'str': 'a', 'int': 1, 'test_label.label': 'no', 'test_label.created': Timestamp(TEST_TIME)},
    {'str': 'b', 'int': 2, 'test_label.label': None, 'test_label.created': None},
    {'str': 'c', 'int': 3, 'test_label.label': None, 'test_label.created': None},
  ]

  assert dataset.get_label_names() == ['test_label']


@freeze_time(TEST_TIME)
def test_add_multiple_labels(make_test_data: TestDataMaker, mocker: MockerFixture) -> None:
  dataset = make_test_data(TEST_ITEMS)

  # Add a 'yes' for every item with int >= 2.
  num_labels = dataset.add_labels('test_label', value='yes', filters=[('int', 'greater_equal', 2)])
  assert num_labels == 2
  # Add a 'no' for every item with int < 2.
  num_labels = dataset.add_labels('test_label', value='no', filters=[('int', 'less', 2)])
  assert num_labels == 1

  assert dataset.manifest() == DatasetManifest(
    source=TestSource(),
    namespace='test_namespace',
    dataset_name='test_dataset',
    data_schema=schema(
      {
        'str': 'string',
        'int': 'int32',
        'test_label': field(fields={'label': 'string', 'created': 'timestamp'}, label='test_label'),
      }
    ),
    num_items=3,
  )

  assert list(dataset.select_rows([PATH_WILDCARD])) == [
    {'str': 'a', 'int': 1, 'test_label.label': 'no', 'test_label.created': TEST_TIME},
    {'str': 'b', 'int': 2, 'test_label.label': 'yes', 'test_label.created': TEST_TIME},
    {'str': 'c', 'int': 3, 'test_label.label': 'yes', 'test_label.created': TEST_TIME},
  ]

  assert dataset.get_label_names() == ['test_label']


@freeze_time(TEST_TIME)
def test_labels_select_groups(make_test_data: TestDataMaker, mocker: MockerFixture) -> None:
  dataset = make_test_data(TEST_ITEMS)

  # Add a 'yes' for every item with int >= 2.
  num_labels = dataset.add_labels('test_label', value='yes', filters=[('int', 'greater_equal', 2)])
  assert num_labels == 2

  # Add a 'no' for every item with int < 2.
  num_labels = dataset.add_labels('test_label', value='no', filters=[('int', 'less', 2)])
  assert num_labels == 1

  assert dataset.select_groups(('test_label', 'label')) == SelectGroupsResult(
    too_many_distinct=False, counts=[('yes', 2), ('no', 1)]
  )

  assert dataset.get_label_names() == ['test_label']
