# Copyright 2021 The Feast Authors
#
# 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.
import logging
import os
import sys
from datetime import timedelta

import pandas as pd
import pytest
from pytest_lazyfixture import lazy_fixture
from testcontainers.core.container import DockerContainer
from testcontainers.core.waiting_utils import wait_for_logs

from feast import FileSource, RequestSource
from feast.data_format import ParquetFormat
from feast.entity import Entity
from feast.errors import FeatureViewNotFoundException
from feast.feature_view import FeatureView
from feast.field import Field
from feast.infra.infra_object import Infra
from feast.infra.online_stores.sqlite import SqliteTable
from feast.infra.registry.sql import SqlRegistry
from feast.on_demand_feature_view import on_demand_feature_view
from feast.repo_config import RegistryConfig
from feast.types import Array, Bytes, Float32, Int32, Int64, String
from feast.value_type import ValueType
from tests.integration.feature_repos.universal.entities import driver

POSTGRES_USER = "test"
POSTGRES_PASSWORD = "test"
POSTGRES_DB = "test"


logger = logging.getLogger(__name__)


@pytest.fixture(scope="session")
def pg_registry():
    container = (
        DockerContainer("postgres:latest")
        .with_exposed_ports(5432)
        .with_env("POSTGRES_USER", POSTGRES_USER)
        .with_env("POSTGRES_PASSWORD", POSTGRES_PASSWORD)
        .with_env("POSTGRES_DB", POSTGRES_DB)
    )

    container.start()

    log_string_to_wait_for = "database system is ready to accept connections"
    waited = wait_for_logs(
        container=container,
        predicate=log_string_to_wait_for,
        timeout=30,
        interval=10,
    )
    logger.info("Waited for %s seconds until postgres container was up", waited)
    container_port = container.get_exposed_port(5432)

    registry_config = RegistryConfig(
        registry_type="sql",
        path=f"postgresql://{POSTGRES_USER}:{POSTGRES_PASSWORD}@127.0.0.1:{container_port}/{POSTGRES_DB}",
    )

    yield SqlRegistry(registry_config, "project", None)

    container.stop()


@pytest.fixture(scope="session")
def mysql_registry():
    container = (
        DockerContainer("mysql:latest")
        .with_exposed_ports(3306)
        .with_env("MYSQL_RANDOM_ROOT_PASSWORD", "true")
        .with_env("MYSQL_USER", POSTGRES_USER)
        .with_env("MYSQL_PASSWORD", POSTGRES_PASSWORD)
        .with_env("MYSQL_DATABASE", POSTGRES_DB)
    )

    container.start()

    # The log string uses '8.0.*' since the version might be changed as new Docker images are pushed.
    log_string_to_wait_for = "/usr/sbin/mysqld: ready for connections. Version: '8.0.*'  socket: '/var/run/mysqld/mysqld.sock'  port: 3306"
    waited = wait_for_logs(
        container=container,
        predicate=log_string_to_wait_for,
        timeout=60,
        interval=10,
    )
    logger.info("Waited for %s seconds until mysql container was up", waited)
    container_port = container.get_exposed_port(3306)

    registry_config = RegistryConfig(
        registry_type="sql",
        path=f"mysql+mysqldb://{POSTGRES_USER}:{POSTGRES_PASSWORD}@127.0.0.1:{container_port}/{POSTGRES_DB}",
    )

    yield SqlRegistry(registry_config, "project", None)

    container.stop()


@pytest.fixture(scope="session")
def sqlite_registry():
    registry_config = RegistryConfig(
        registry_type="sql",
        path="sqlite://",
    )

    yield SqlRegistry(registry_config, "project", None)


@pytest.mark.skipif(
    sys.platform == "darwin" and "GITHUB_REF" in os.environ,
    reason="does not run on mac github actions",
)
@pytest.mark.parametrize(
    "sql_registry",
    [
        lazy_fixture("mysql_registry"),
        lazy_fixture("pg_registry"),
        lazy_fixture("sqlite_registry"),
    ],
)
def test_apply_entity_success(sql_registry):
    entity = Entity(
        name="driver_car_id",
        description="Car driver id",
        tags={"team": "matchmaking"},
    )

    project = "project"

    # Register Entity
    sql_registry.apply_entity(entity, project)
    project_metadata = sql_registry.list_project_metadata(project=project)
    assert len(project_metadata) == 1
    project_uuid = project_metadata[0].project_uuid
    assert len(project_metadata[0].project_uuid) == 36
    assert_project_uuid(project, project_uuid, sql_registry)

    entities = sql_registry.list_entities(project)
    assert_project_uuid(project, project_uuid, sql_registry)

    entity = entities[0]
    assert (
        len(entities) == 1
        and entity.name == "driver_car_id"
        and entity.description == "Car driver id"
        and "team" in entity.tags
        and entity.tags["team"] == "matchmaking"
    )

    entity = sql_registry.get_entity("driver_car_id", project)
    assert (
        entity.name == "driver_car_id"
        and entity.description == "Car driver id"
        and "team" in entity.tags
        and entity.tags["team"] == "matchmaking"
    )

    # After the first apply, the created_timestamp should be the same as the last_update_timestamp.
    assert entity.created_timestamp == entity.last_updated_timestamp

    sql_registry.delete_entity("driver_car_id", project)
    assert_project_uuid(project, project_uuid, sql_registry)
    entities = sql_registry.list_entities(project)
    assert_project_uuid(project, project_uuid, sql_registry)
    assert len(entities) == 0

    sql_registry.teardown()


def assert_project_uuid(project, project_uuid, sql_registry):
    project_metadata = sql_registry.list_project_metadata(project=project)
    assert len(project_metadata) == 1
    assert project_metadata[0].project_uuid == project_uuid


@pytest.mark.skipif(
    sys.platform == "darwin" and "GITHUB_REF" in os.environ,
    reason="does not run on mac github actions",
)
@pytest.mark.parametrize(
    "sql_registry",
    [
        lazy_fixture("mysql_registry"),
        lazy_fixture("pg_registry"),
        lazy_fixture("sqlite_registry"),
    ],
)
def test_apply_feature_view_success(sql_registry):
    # Create Feature Views
    batch_source = FileSource(
        file_format=ParquetFormat(),
        path="file://feast/*",
        timestamp_field="ts_col",
        created_timestamp_column="timestamp",
    )

    entity = Entity(name="fs1_my_entity_1", join_keys=["test"])

    fv1 = FeatureView(
        name="my_feature_view_1",
        schema=[
            Field(name="fs1_my_feature_1", dtype=Int64),
            Field(name="fs1_my_feature_2", dtype=String),
            Field(name="fs1_my_feature_3", dtype=Array(String)),
            Field(name="fs1_my_feature_4", dtype=Array(Bytes)),
        ],
        entities=[entity],
        tags={"team": "matchmaking"},
        source=batch_source,
        ttl=timedelta(minutes=5),
    )

    project = "project"

    # Register Feature View
    sql_registry.apply_feature_view(fv1, project)

    feature_views = sql_registry.list_feature_views(project)

    # List Feature Views
    assert (
        len(feature_views) == 1
        and feature_views[0].name == "my_feature_view_1"
        and feature_views[0].features[0].name == "fs1_my_feature_1"
        and feature_views[0].features[0].dtype == Int64
        and feature_views[0].features[1].name == "fs1_my_feature_2"
        and feature_views[0].features[1].dtype == String
        and feature_views[0].features[2].name == "fs1_my_feature_3"
        and feature_views[0].features[2].dtype == Array(String)
        and feature_views[0].features[3].name == "fs1_my_feature_4"
        and feature_views[0].features[3].dtype == Array(Bytes)
        and feature_views[0].entities[0] == "fs1_my_entity_1"
    )

    feature_view = sql_registry.get_feature_view("my_feature_view_1", project)
    assert (
        feature_view.name == "my_feature_view_1"
        and feature_view.features[0].name == "fs1_my_feature_1"
        and feature_view.features[0].dtype == Int64
        and feature_view.features[1].name == "fs1_my_feature_2"
        and feature_view.features[1].dtype == String
        and feature_view.features[2].name == "fs1_my_feature_3"
        and feature_view.features[2].dtype == Array(String)
        and feature_view.features[3].name == "fs1_my_feature_4"
        and feature_view.features[3].dtype == Array(Bytes)
        and feature_view.entities[0] == "fs1_my_entity_1"
    )
    assert feature_view.ttl == timedelta(minutes=5)

    # After the first apply, the created_timestamp should be the same as the last_update_timestamp.
    assert feature_view.created_timestamp == feature_view.last_updated_timestamp

    # Modify the feature view and apply again to test if diffing the online store table works
    fv1.ttl = timedelta(minutes=6)
    sql_registry.apply_feature_view(fv1, project)
    feature_views = sql_registry.list_feature_views(project)
    assert len(feature_views) == 1
    feature_view = sql_registry.get_feature_view("my_feature_view_1", project)
    assert feature_view.ttl == timedelta(minutes=6)

    # Delete feature view
    sql_registry.delete_feature_view("my_feature_view_1", project)
    feature_views = sql_registry.list_feature_views(project)
    assert len(feature_views) == 0

    sql_registry.teardown()


@pytest.mark.skipif(
    sys.platform == "darwin" and "GITHUB_REF" in os.environ,
    reason="does not run on mac github actions",
)
@pytest.mark.parametrize(
    "sql_registry",
    [
        lazy_fixture("mysql_registry"),
        lazy_fixture("pg_registry"),
        lazy_fixture("sqlite_registry"),
    ],
)
def test_apply_on_demand_feature_view_success(sql_registry):
    # Create Feature Views
    driver_stats = FileSource(
        name="driver_stats_source",
        path="data/driver_stats_lat_lon.parquet",
        timestamp_field="event_timestamp",
        created_timestamp_column="created",
        description="A table describing the stats of a driver based on hourly logs",
        owner="test2@gmail.com",
    )

    driver_daily_features_view = FeatureView(
        name="driver_daily_features",
        entities=[driver()],
        ttl=timedelta(seconds=8640000000),
        schema=[
            Field(name="daily_miles_driven", dtype=Float32),
            Field(name="lat", dtype=Float32),
            Field(name="lon", dtype=Float32),
            Field(name="string_feature", dtype=String),
        ],
        online=True,
        source=driver_stats,
        tags={"production": "True"},
        owner="test2@gmail.com",
    )

    @on_demand_feature_view(
        sources=[driver_daily_features_view],
        schema=[Field(name="first_char", dtype=String)],
    )
    def location_features_from_push(inputs: pd.DataFrame) -> pd.DataFrame:
        df = pd.DataFrame()
        df["first_char"] = inputs["string_feature"].str[:1].astype("string")
        return df

    project = "project"

    with pytest.raises(FeatureViewNotFoundException):
        sql_registry.get_user_metadata(project, location_features_from_push)

    # Register Feature View
    sql_registry.apply_feature_view(location_features_from_push, project)

    assert not sql_registry.get_user_metadata(project, location_features_from_push)

    b = "metadata".encode("utf-8")
    sql_registry.apply_user_metadata(project, location_features_from_push, b)
    assert sql_registry.get_user_metadata(project, location_features_from_push) == b

    feature_views = sql_registry.list_on_demand_feature_views(project)

    # List Feature Views
    assert (
        len(feature_views) == 1
        and feature_views[0].name == "location_features_from_push"
        and feature_views[0].features[0].name == "first_char"
        and feature_views[0].features[0].dtype == String
    )

    feature_view = sql_registry.get_on_demand_feature_view(
        "location_features_from_push", project
    )
    assert (
        feature_view.name == "location_features_from_push"
        and feature_view.features[0].name == "first_char"
        and feature_view.features[0].dtype == String
    )

    sql_registry.delete_feature_view("location_features_from_push", project)
    feature_views = sql_registry.list_on_demand_feature_views(project)
    assert len(feature_views) == 0

    sql_registry.teardown()


@pytest.mark.skipif(
    sys.platform == "darwin" and "GITHUB_REF" in os.environ,
    reason="does not run on mac github actions",
)
@pytest.mark.parametrize(
    "sql_registry",
    [
        lazy_fixture("mysql_registry"),
        lazy_fixture("pg_registry"),
        lazy_fixture("sqlite_registry"),
    ],
)
def test_modify_feature_views_success(sql_registry):
    # Create Feature Views
    batch_source = FileSource(
        file_format=ParquetFormat(),
        path="file://feast/*",
        timestamp_field="ts_col",
        created_timestamp_column="timestamp",
    )

    request_source = RequestSource(
        name="request_source",
        schema=[Field(name="my_input_1", dtype=Int32)],
    )

    entity = Entity(name="fs1_my_entity_1", join_keys=["test"])

    fv1 = FeatureView(
        name="my_feature_view_1",
        schema=[Field(name="fs1_my_feature_1", dtype=Int64)],
        entities=[entity],
        tags={"team": "matchmaking"},
        source=batch_source,
        ttl=timedelta(minutes=5),
    )

    @on_demand_feature_view(
        schema=[
            Field(name="odfv1_my_feature_1", dtype=String),
            Field(name="odfv1_my_feature_2", dtype=Int32),
        ],
        sources=[request_source],
    )
    def odfv1(feature_df: pd.DataFrame) -> pd.DataFrame:
        data = pd.DataFrame()
        data["odfv1_my_feature_1"] = feature_df["my_input_1"].astype("category")
        data["odfv1_my_feature_2"] = feature_df["my_input_1"].astype("int32")
        return data

    project = "project"

    # Register Feature Views
    sql_registry.apply_feature_view(odfv1, project)
    sql_registry.apply_feature_view(fv1, project)

    # Modify odfv by changing a single feature dtype
    @on_demand_feature_view(
        schema=[
            Field(name="odfv1_my_feature_1", dtype=Float32),
            Field(name="odfv1_my_feature_2", dtype=Int32),
        ],
        sources=[request_source],
    )
    def odfv1(feature_df: pd.DataFrame) -> pd.DataFrame:
        data = pd.DataFrame()
        data["odfv1_my_feature_1"] = feature_df["my_input_1"].astype("float")
        data["odfv1_my_feature_2"] = feature_df["my_input_1"].astype("int32")
        return data

    # Apply the modified odfv
    sql_registry.apply_feature_view(odfv1, project)

    # Check odfv
    on_demand_feature_views = sql_registry.list_on_demand_feature_views(project)

    assert (
        len(on_demand_feature_views) == 1
        and on_demand_feature_views[0].name == "odfv1"
        and on_demand_feature_views[0].features[0].name == "odfv1_my_feature_1"
        and on_demand_feature_views[0].features[0].dtype == Float32
        and on_demand_feature_views[0].features[1].name == "odfv1_my_feature_2"
        and on_demand_feature_views[0].features[1].dtype == Int32
    )
    request_schema = on_demand_feature_views[0].get_request_data_schema()
    assert (
        list(request_schema.keys())[0] == "my_input_1"
        and list(request_schema.values())[0] == ValueType.INT32
    )

    feature_view = sql_registry.get_on_demand_feature_view("odfv1", project)
    assert (
        feature_view.name == "odfv1"
        and feature_view.features[0].name == "odfv1_my_feature_1"
        and feature_view.features[0].dtype == Float32
        and feature_view.features[1].name == "odfv1_my_feature_2"
        and feature_view.features[1].dtype == Int32
    )
    request_schema = feature_view.get_request_data_schema()
    assert (
        list(request_schema.keys())[0] == "my_input_1"
        and list(request_schema.values())[0] == ValueType.INT32
    )

    # Make sure fv1 is untouched
    feature_views = sql_registry.list_feature_views(project)

    # List Feature Views
    assert (
        len(feature_views) == 1
        and feature_views[0].name == "my_feature_view_1"
        and feature_views[0].features[0].name == "fs1_my_feature_1"
        and feature_views[0].features[0].dtype == Int64
        and feature_views[0].entities[0] == "fs1_my_entity_1"
    )

    feature_view = sql_registry.get_feature_view("my_feature_view_1", project)
    assert (
        feature_view.name == "my_feature_view_1"
        and feature_view.features[0].name == "fs1_my_feature_1"
        and feature_view.features[0].dtype == Int64
        and feature_view.entities[0] == "fs1_my_entity_1"
    )

    sql_registry.teardown()


@pytest.mark.skipif(
    sys.platform == "darwin" and "GITHUB_REF" in os.environ,
    reason="does not run on mac github actions",
)
@pytest.mark.parametrize(
    "sql_registry",
    [
        lazy_fixture("mysql_registry"),
        lazy_fixture("pg_registry"),
        lazy_fixture("sqlite_registry"),
    ],
)
def test_apply_data_source(sql_registry):
    # Create Feature Views
    batch_source = FileSource(
        name="test_source",
        file_format=ParquetFormat(),
        path="file://feast/*",
        timestamp_field="ts_col",
        created_timestamp_column="timestamp",
    )

    entity = Entity(name="fs1_my_entity_1", join_keys=["test"])

    fv1 = FeatureView(
        name="my_feature_view_1",
        schema=[
            Field(name="fs1_my_feature_1", dtype=Int64),
            Field(name="fs1_my_feature_2", dtype=String),
            Field(name="fs1_my_feature_3", dtype=Array(String)),
            Field(name="fs1_my_feature_4", dtype=Array(Bytes)),
        ],
        entities=[entity],
        tags={"team": "matchmaking"},
        source=batch_source,
        ttl=timedelta(minutes=5),
    )

    project = "project"

    # Register data source and feature view
    sql_registry.apply_data_source(batch_source, project, commit=False)
    sql_registry.apply_feature_view(fv1, project, commit=True)

    registry_feature_views = sql_registry.list_feature_views(project)
    registry_data_sources = sql_registry.list_data_sources(project)
    assert len(registry_feature_views) == 1
    assert len(registry_data_sources) == 1
    registry_feature_view = registry_feature_views[0]
    assert registry_feature_view.batch_source == batch_source
    registry_data_source = registry_data_sources[0]
    assert registry_data_source == batch_source

    # Check that change to batch source propagates
    batch_source.timestamp_field = "new_ts_col"
    sql_registry.apply_data_source(batch_source, project, commit=False)
    sql_registry.apply_feature_view(fv1, project, commit=True)
    registry_feature_views = sql_registry.list_feature_views(project)
    registry_data_sources = sql_registry.list_data_sources(project)
    assert len(registry_feature_views) == 1
    assert len(registry_data_sources) == 1
    registry_feature_view = registry_feature_views[0]
    assert registry_feature_view.batch_source == batch_source
    registry_batch_source = sql_registry.list_data_sources(project)[0]
    assert registry_batch_source == batch_source

    sql_registry.teardown()


@pytest.mark.skipif(
    sys.platform == "darwin" and "GITHUB_REF" in os.environ,
    reason="does not run on mac github actions",
)
@pytest.mark.parametrize(
    "sql_registry",
    [
        lazy_fixture("mysql_registry"),
        lazy_fixture("pg_registry"),
        lazy_fixture("sqlite_registry"),
    ],
)
def test_registry_cache(sql_registry):
    # Create Feature Views
    batch_source = FileSource(
        name="test_source",
        file_format=ParquetFormat(),
        path="file://feast/*",
        timestamp_field="ts_col",
        created_timestamp_column="timestamp",
    )

    entity = Entity(name="fs1_my_entity_1", join_keys=["test"])

    fv1 = FeatureView(
        name="my_feature_view_1",
        schema=[
            Field(name="fs1_my_feature_1", dtype=Int64),
            Field(name="fs1_my_feature_2", dtype=String),
            Field(name="fs1_my_feature_3", dtype=Array(String)),
            Field(name="fs1_my_feature_4", dtype=Array(Bytes)),
        ],
        entities=[entity],
        tags={"team": "matchmaking"},
        source=batch_source,
        ttl=timedelta(minutes=5),
    )

    project = "project"

    # Register data source and feature view
    sql_registry.apply_data_source(batch_source, project)
    sql_registry.apply_feature_view(fv1, project)
    registry_feature_views_cached = sql_registry.list_feature_views(
        project, allow_cache=True
    )
    registry_data_sources_cached = sql_registry.list_data_sources(
        project, allow_cache=True
    )
    # Not refreshed cache, so cache miss
    assert len(registry_feature_views_cached) == 0
    assert len(registry_data_sources_cached) == 0
    sql_registry.refresh(project)
    # Now objects exist
    registry_feature_views_cached = sql_registry.list_feature_views(
        project, allow_cache=True
    )
    registry_data_sources_cached = sql_registry.list_data_sources(
        project, allow_cache=True
    )
    assert len(registry_feature_views_cached) == 1
    assert len(registry_data_sources_cached) == 1
    registry_feature_view = registry_feature_views_cached[0]
    assert registry_feature_view.batch_source == batch_source
    registry_data_source = registry_data_sources_cached[0]
    assert registry_data_source == batch_source

    sql_registry.teardown()


@pytest.mark.skipif(
    sys.platform == "darwin" and "GITHUB_REF" in os.environ,
    reason="does not run on mac github actions",
)
@pytest.mark.parametrize(
    "sql_registry",
    [
        lazy_fixture("mysql_registry"),
        lazy_fixture("pg_registry"),
        lazy_fixture("sqlite_registry"),
    ],
)
def test_update_infra(sql_registry):
    # Create infra object
    project = "project"
    infra = sql_registry.get_infra(project=project)

    assert len(infra.infra_objects) == 0

    # Should run update infra successfully
    sql_registry.update_infra(infra, project)

    # Should run update infra successfully when adding
    new_infra = Infra()
    new_infra.infra_objects.append(
        SqliteTable(
            path="/tmp/my_path.db",
            name="my_table",
        )
    )
    sql_registry.update_infra(new_infra, project)
    infra = sql_registry.get_infra(project=project)
    assert len(infra.infra_objects) == 1

    # Try again since second time, infra should be not-empty
    sql_registry.teardown()