# -*- coding: utf-8 -*-
import os
import argparse

from datacloud.dev_utils.yt import yt_utils
from datacloud.dev_utils.logging.logger import get_basic_logger
from datacloud.dev_utils.yql.yql_helpers import create_yql_client
from datacloud.stability.score_acceptance.assert_stability import DEFAULT_AUC_THRESH, DEFAULT_HIT_THRESH
from datacloud.cli.upload_model.helpers import (
    prepare_batch, upload_binary, add_config_to_yt,
    add_score_path_to_ydb, add_score_to_internal_ydb
)

logger = get_basic_logger(__name__)

MODEL_TYPES_CHOICE = ('PredictModel', 'PredictProbaModel', 'InvertedPredictProbaModel',
                      'CatboostPredictProbaModel')
FEATURES_CHOICE = ('DSSMFeatureBase400', 'ClusterFeatureBase763', 'ClusterFeatureBase767', 'NormedS2VFeatureBase512', 'ContactActionFeature', 'TimeHistFeature', 'PhoneRangeFeature')


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-p', '--partner-id', help='<Required> Partner ID', required=True)
    parser.add_argument('-s', '--score-name', help='<Required> Score name', required=True)
    parser.add_argument('-l', '--local-binary', help='<Required> Path to local binary', required=True)
    parser.add_argument('-c', '--model-class', help='<Required> Model class',
                        choices=MODEL_TYPES_CHOICE, required=True)
    parser.add_argument('-f', '--features', help='<Required> Features', nargs='+',
                        choices=FEATURES_CHOICE, required=True)
    parser.add_argument('-r', '--path-to-raw', help='Path to raw table for batch (NOT required)',
                        required=False)
    parser.add_argument('-a', '--auc-thresh', help='AUC threshold for batch monitoring',
                        default=DEFAULT_AUC_THRESH, required=False)
    parser.add_argument('-i', '--hit-thresh', help='Hit threshold for batch monitoring',
                        default=DEFAULT_HIT_THRESH, required=False)
    parser.add_argument('-t', '--target', help='Target field name for batch monitoring',
                        default='target', required=False)
    parser.add_argument('-y', '--yt-binary-name', help='YT binary name', required=False)
    parser.add_argument('--not-active', help='Set model not active', action='store_true')
    parser.add_argument('--save-info', help='Save model info (features, etc)', action='store_true')
    parser.add_argument('--cookie-sync-on', help='Whether score requires cookie sync',
                        action='store_true')
    parser.add_argument('--transfer-off', help='Whether to off regular transfer to YT',
                        action='store_true')

    return parser.parse_args()


def main():
    args = get_args()
    yt_client = yt_utils.get_yt_client()

    yql_token = os.environ.get('YQL_TOKEN')
    if yql_token is not None:
        yql_client = create_yql_client(yt_client=yt_client, token=yql_token)
        logger.info('Got token from YQL_TOKEN env')
    else:
        yql_client = create_yql_client(yt_client=yt_client)
        logger.info('Got token from yt config')
    assert yql_token, 'Please provide YQL token'

    if args.path_to_raw:
        prepare_batch(yt_client, args)
    upload_binary(yt_client, args)
    add_config_to_yt(yt_client, args)

    configure_ydb = False
    if configure_ydb:
        add_score_path_to_ydb(yql_client, args)
        add_score_to_internal_ydb(yql_token, args)


if __name__ == '__main__':
    main()
