#
# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
#

import pytest

from snowflake.core.cortex.search_service import QueryRequest


TEST_SERVICE_NAME = "SNOWPY_TEST_SERVICE"

API_NOT_ENABLED = 'Cortex Search API is not enabled'

@pytest.fixture(scope='session')
@pytest.mark.usefixtures("backup_warehouse")
def setup_cortex_search_service(connection, database, schema, warehouse):
    _database_name = database.name
    _schema_name = schema.name
    with connection.cursor() as cursor:
        cursor.execute(f"use warehouse {warehouse.name}").fetchone()

        test_table_name = f"{_database_name}.{_schema_name}.SNOWPY_TEST_TABLE"
        # Base Table
        cursor.execute(
            f"CREATE OR REPLACE TABLE {test_table_name} (col1 VARCHAR, col2 VARCHAR)",
        )

        rows = ",".join(["('hi', 'hello')"] * 20)
        cursor.execute(
            f"INSERT INTO {test_table_name} VALUES {rows}",
        )

        # Cortex Search Service
        cursor.execute(
            f"CREATE OR REPLACE CORTEX SEARCH SERVICE {_database_name}.{_schema_name}.{TEST_SERVICE_NAME} "
            f"ON col1 TARGET_LAG='1 minute' WAREHOUSE={warehouse.name} "
            f"AS SELECT col1, col2 FROM {test_table_name}",
        )

        try:
            yield
        finally:
            cursor.execute(f"DROP CORTEX SEARCH SERVICE {_database_name}.{_schema_name}.{TEST_SERVICE_NAME}")
            cursor.execute(f"DROP TABLE {test_table_name}")

@pytest.fixture(scope='session')
@pytest.mark.usefixtures('setup_cortex_search_service')
def precheck_cortex_search_enabled(cortex_search_services):
    try:
        cortex_search_services[TEST_SERVICE_NAME].search("hi", ["col1", "col2"], limit=5)
    except Exception as err:
        if API_NOT_ENABLED in err.body:
            pytest.xfail(API_NOT_ENABLED)
        raise err

pytestmark = pytest.mark.usefixtures("setup_cortex_search_service", "precheck_cortex_search_enabled")

def test_search(cortex_search_services):
    resp = cortex_search_services[TEST_SERVICE_NAME].search("hi", ["col1", "col2"], limit=5)
    assert len(resp.results) == 5
    for row in resp.results:
        assert row["col1"] is not None
        assert row["col2"] is not None


def test_search_collection(cortex_search_services):
    resp = cortex_search_services.search(
        TEST_SERVICE_NAME,
        QueryRequest.from_dict({"query": "hi", "columns": ["col1", "col2"], "limit": 5}),
    )

    assert len(resp.results) == 5
    for row in resp.results:
        assert row["col1"] is not None
        assert row["col2"] is not None
