import pytest
import workbench  # noqa: F401
import logging
from workbench.utils.test_data_generator import TestDataGenerator
from workbench.api import DataSource, FeatureSet, Model, Endpoint, ModelType

# Set the logging level
logging.getLogger("workbench").setLevel(logging.DEBUG)


def create_data_source():
    test_data = TestDataGenerator()
    df = test_data.person_data()
    if not DataSource("abc_test").ready():
        DataSource(df, name="abc_test")


def create_feature_set():
    create_data_source()

    # If the feature set doesn't exist, create it
    if not FeatureSet("abc_features").ready():
        DataSource("abc_test").to_features("abc_features", id_column="id")


def create_model():
    create_feature_set()

    # If the model doesn't exist, create it
    if not Model("abc-regression").ready():
        FeatureSet("abc_features").to_model(
            name="abc-regression", model_type=ModelType.REGRESSOR, target_column="iq_score"
        )


def create_endpoint():
    create_model()

    # Create some new endpoints
    if not Endpoint("abc-regression").ready():
        Model("abc-regression").to_endpoint(name="abc-regression")


@pytest.mark.long
def test_endpoint_deletion():
    create_endpoint()

    # Now Delete the endpoint
    Endpoint("abc-regression").delete()


@pytest.mark.long
def test_model_deletion():
    create_model()

    # Now Delete the Model
    Model("abc-regression").delete()


@pytest.mark.long
def test_feature_set_deletion():
    create_feature_set()

    # Now Delete the FeatureSet
    FeatureSet("abc_features").delete()


@pytest.mark.long
def test_data_source_deletion():
    create_data_source()

    # Now Delete the DataSource
    DataSource("abc_test").delete()


if __name__ == "__main__":

    test_endpoint_deletion()
    test_model_deletion()
    test_feature_set_deletion()
    test_data_source_deletion()
