import os
from yt.wrapper import ypath_join

from datacloud.dev_utils.transfer.yt_to_ydb.score_path_table import ScorePathTableDescription
from datacloud.dev_utils.yt import yt_files
from datacloud.config.yt import PARTNER_BINARY_MODEL_YT_PATH
from datacloud.dev_utils.transfer.yt_to_ydb.ydb_constants import (
    YDB_ENDPOINT, YDB_DATABASE, YDB_ROOT
)
from datacloud.model_applyer.tables.models_config_table import ApiModelsConfigTable
from datacloud.dev_utils.transfer.yt_to_ydb import table_description
from datacloud.dev_utils.logging.logger import get_basic_logger
from datacloud.dev_utils.ydb.lib.core.ydb_manager import YdbManager, YdbConnectionParams
from datacloud.score_api.storage.ydb.ydb_tables.config_tables.partner_scores_table import PartnerScoresTable
from datacloud.stability.mini_batch_monitoring.perform import INTERNAL_PARTNER
from datacloud.stability.batch_monitoring import TEST_SAMPLES_FOLDER

from datacloud.cli.upload_model.mr import BatchMapper

logger = get_basic_logger(__name__)

YDB_SCORE_PATH = os.path.join(YDB_ROOT, 'config', 'score_path')
YDB_PARTNER_SCORES = os.path.join(YDB_ROOT, 'config', 'partner_scores')


def prepare_batch(yt_client, args):
    logger.info('Preparing batch...')
    path_to_batch = ypath_join(TEST_SAMPLES_FOLDER, args.score_name)
    assert not yt_client.exists(path_to_batch), 'Batch already exists!'
    col_names = set(col['name'] for col in yt_client.get_attribute(args.path_to_raw, 'schema'))
    assert args.target in col_names, 'Column "{}" not found in schema!'.format(args.target)

    batch_table = yt_client.TablePath(
        path_to_batch,
        schema=[
            {'name': 'id_value', 'type': 'string'},
            {'name': 'external_id', 'type': 'string'},
            {'name': 'target', 'type': 'int64'},
        ]
    )

    with yt_client.Transaction():
        yt_client.run_map(BatchMapper(args.target), args.path_to_raw, batch_table,
                          spec={'title': 'Preparing batch for {}'.format(args.score_name)})
        assert yt_client.row_count(batch_table) > 0, 'Something went wrong! Batch table is empty!'

        yt_client.run_sort(
            batch_table,
            sort_by='id_value',
            spec={'title': 'Preparing batch for {} / sort after'.format(args.score_name)})

    logger.info('Batch prepared!')


def upload_binary(yt_client, args):
    logger.info('Uploading binary to YT...')
    yt_binary_name = args.yt_binary_name or os.path.split(args.local_binary)[1]
    yt_root = ypath_join(PARTNER_BINARY_MODEL_YT_PATH, args.partner_id, args.score_name)

    yt_binary_path = ypath_join(yt_root, yt_binary_name)
    yt_files.upload_file(yt_client, args.local_binary, yt_binary_path)
    logger.info('Binary uploaded to YT!')

    logger.info('Setting link to binary on YT...')
    yt_link_path = ypath_join(yt_root, 'current')
    yt_client.link(yt_binary_path, yt_link_path, force=True)
    logger.info('Link to binary set on YT!')


def add_config_to_yt(yt_client, args):
    logger.info('Adding model to Api Models Config Table...')
    features = ' '.join(args.features)
    additional = {'transfer_on': not args.transfer_off}
    if args.cookie_sync_on:
        additional['cookie_sync_on'] = True
    if args.path_to_raw:
        additional['batch_monitoring'] = [{
            'batch_name': args.score_name,
            'AUC_threshold': float(args.auc_thresh),
            'hit_threshold': float(args.hit_thresh),
        }]
    ApiModelsConfigTable().add_model(partner_id=args.partner_id, score_name=args.score_name,
                                     model_class=args.model_class, features=features,
                                     is_active=(not args.not_active), save_info=args.save_info,
                                     additional=additional)
    logger.info('Model added to Api Models Config Table!')


def add_score_path_to_ydb(yql_client, args):
    logger.info('Adding path to score path YDB table...')
    ydb_connection_params = table_description.YdbConnectionParams(
        endpoint=YDB_ENDPOINT,
        database=YDB_DATABASE
    )
    score_path_table_description = ScorePathTableDescription(YDB_SCORE_PATH, ydb_connection_params)
    score_path_table_description.insert(yql_client, [score_path_table_description.Record(
        internal_score_name=args.score_name
    )])
    logger.info('Path added to score path YDB table!')


def add_score_to_internal_ydb(yql_token, args):
    logger.info('Adding access to internal partner...')
    connection_params = YdbConnectionParams(
        endpoint=YDB_ENDPOINT, database=YDB_DATABASE, auth_token=yql_token)
    partner_scores_table = PartnerScoresTable(
        ydb_manager=YdbManager(connection_params), database=YDB_DATABASE,
        table=YDB_PARTNER_SCORES)

    partner_scores_table.insert([PartnerScoresTable.Record(
        partner_id=INTERNAL_PARTNER,
        partner_score_name=args.score_name,
        internal_score_name=args.score_name,
        is_active=True
    )])
    logger.info('Access to internal provided!')
