from pathlib import Path

from tadas.cli import launcher
from tadas.domain.constants import REPORT_FINAL_TABLES
from tadas.snowflake import snowflake_publish
from tadas.snowflake import client as snowflake_utils
from tadas.platform import locking as lock_utils
from tadas.domain import constants

THIS_DIR = Path(__file__).parent.resolve()
LOCK_PATH = THIS_DIR / 'common_test_models.lock'

EXPECTED_COLUMNS = snowflake_publish.EXPECTED_COLUMNS

COLUMNS_WITH_VALUES = [
    column
    for column in EXPECTED_COLUMNS if column not in ['report_date', 'geo_country']
]

EXPECTED_TRENDING_FLAGS = snowflake_publish.TRENDING_FLAGS_COLUMNS

EXPECTED_COUNTRIES = sorted(constants.geos + ['XX'])

# these some small countries can have 0 trending tracks
COUNTRIES_OK_IF_EMPTY = {
    'BM',
    'BS',
    'DM',
    'GD',
    'HK',
    'HT',
    'JM',
    'MS',
    'PA',
    'UY',
    'VE',
    'VN',
}


def clear_final_tables():
    with snowflake_utils.snowflake_connection() as conn:
        for final_table in REPORT_FINAL_TABLES.values():
            snowflake_utils.execute(conn, f'DELETE FROM {final_table}')

def run_model(model_version, report_date):
    """
    Run the given model versions for a given report_date.

    :param model_version: model version to run
    :param report_date: report_date in YYYY-MM-DD format
    :return: None
    """
    targets = ['clean', 'inference', 'publish']
    launcher.run_makefile(
        tadas_model_version=model_version,
        args=targets,
        env={
            'REPORT_DATE': report_date,
        },
    )


def lock_final_tables():
    """
    :return: This lock should be used in every integration test that modifies final tables.
    """
    return lock_utils.try_lock(LOCK_PATH)