# -*- coding: utf-8 -*-
import os
import yt.wrapper as yt_wrapper
from yt.wrapper import ypath_split

from datacloud.dev_utils.yt.yt_utils import get_yt_client
from datacloud.dev_utils.logging.logger import get_basic_logger
from datacloud.dev_utils.transfer.yt_to_ydb import table_description, transfer_table
from datacloud.dev_utils.yql import yql_helpers
from datacloud.dev_utils.time.utils import now_str
from datacloud.dev_utils.transfer.yt_to_ydb.geo_table_descriptions import YdbGeoLogsTableDescription, \
    InterestingCryptaTableDescription, GeoTablesListTableDescription, GeoPathTableDescription
from datacloud.launcher.lib.routines.geo_logs import INTERESTING_CRYPTA_SUFFIX, GeoLogsAggregator
from datacloud.dev_utils.status_db.db import lock_manager, LockError
from datacloud.dev_utils.transfer.yt_to_ydb.ydb_constants import (
    YDB_ENDPOINT, YDB_DATABASE, YDB_ROOT, YDB_LOCK_NAME
)
from datacloud.dev_utils.status_db.task import Task, Status

logger = get_basic_logger(__name__)

YDB_GEO_ROOT = os.path.join(YDB_ROOT, 'geo')
YDB_GEO_PATH_CONFIG = os.path.join(YDB_ROOT, 'config', 'geo_path')
YDB_GEO_TABLES_LIST_TABLE = os.path.join(YDB_ROOT, 'config', 'geo_tables_list')

MAX_GEO_TABLES = 2


def get_yql_token():
    yql_token = os.environ.get('YQL_TOKEN')
    if not yql_token:
        yql_token = yt_wrapper.config['token']
    assert yql_token, 'Please provide YQL token'

    return yql_token


def transfer_geo_to_ydb(task, is_transfer_crypta):
    yql_token = get_yql_token()
    logger.info(' Started transfering %s to YDB', task)
    input_path = task.data['table_path']

    yt_client = get_yt_client()
    geo_agg = GeoLogsAggregator(yt_client)

    if is_transfer_crypta:
        local_table = geo_agg.get_interesting_crypta_table(input_path)
    else:
        local_table = geo_agg.get_prod_aggregate_table(input_path)

    yql_client = yql_helpers.create_yql_client(yt_client, token=yql_token)

    ydb_connection_params = table_description.YdbConnectionParams(
        endpoint=YDB_ENDPOINT,
        database=YDB_DATABASE
    )

    source_table = table_description.YtTableDescription(
        table_path=local_table,
        yt_cluster=yt_wrapper.config['proxy']['url']
    )

    if is_transfer_crypta:
        TargetTableClass = InterestingCryptaTableDescription
    else:
        TargetTableClass = YdbGeoLogsTableDescription

    target_table = TargetTableClass(
        table_path=os.path.join(YDB_GEO_ROOT, ypath_split(local_table)[-1]),
        ydb_connection_params=ydb_connection_params
    )

    try:
        with lock_manager(YDB_LOCK_NAME):
            transfer_table.create_table(yql_client, target_table)
            transfer_table.transfer(yql_client, source_table, target_table, yql_token)
    except LockError:
        logger.error('Transfer is loced!')
        return []

    current_time = now_str()
    if is_transfer_crypta:
        next_program = Task('run_rotate_geo_tables', input_path, Status.READY, {'table_path': input_path}, current_time, current_time)
    else:
        next_program = Task('transfer_crypta_to_ydb', input_path, Status.READY, {'table_path': input_path}, current_time, current_time)

    new_programs = [
        task.make_done(),
        next_program
    ]
    return new_programs


def rotate_geo_tables(task):
    input_path = task.data['table_path']
    yql_token = get_yql_token()

    logger.info(' Switching tables at YDB %s', input_path)
    date = input_path.split(':')[0]

    yql_client = yql_helpers.create_yql_client(db='[{}]'.format(YDB_DATABASE), token=yql_token)

    ydb_connection_params = table_description.YdbConnectionParams(
        endpoint=YDB_ENDPOINT,
        database=YDB_DATABASE
    )

    geo_tables_list_table = GeoTablesListTableDescription(
        table_path=YDB_GEO_TABLES_LIST_TABLE,
        ydb_connection_params=ydb_connection_params
    )

    geo_tables_list = geo_tables_list_table.select(yql_client)
    assert len(geo_tables_list) <= MAX_GEO_TABLES, 'Too much ({} > {}) backups!'.format(
        len(geo_tables_list),
        MAX_GEO_TABLES
    )

    new_tables_list_record = GeoTablesListTableDescription.Record(
        date=date,
        crypta=input_path + INTERESTING_CRYPTA_SUFFIX,
        log=input_path,
        is_delete_permited=True,
        additional=''
    )
    geo_tables_list_table.insert(yql_client, new_tables_list_record)

    path_table = GeoPathTableDescription(
        table_path=YDB_GEO_PATH_CONFIG,
        ydb_connection_params=ydb_connection_params
    )

    if len(geo_tables_list) > 0:
        newest_crypta = geo_tables_list[-1].crypta
        newest_geo = geo_tables_list[-1].log

        path_table.replace_paths(
            yql_client,
            old_crypta_path=newest_crypta,
            old_geo_path=newest_geo,
            new_crypta_path=input_path + INTERESTING_CRYPTA_SUFFIX,
            new_geo_path=input_path
        )

    if len(geo_tables_list) == MAX_GEO_TABLES:
        oldest_crypta = geo_tables_list[0].crypta
        oldest_geo = geo_tables_list[0].log
        oldest_date = geo_tables_list[0].date

        oldest_records_in_geo_path = path_table.select(
            yql_client=yql_client,
            crypta_path=oldest_crypta,
            geo_path=oldest_geo
        )

        if oldest_records_in_geo_path:
            raise RuntimeError('Some tables to delete are used in geo_path!')

        geo_tables_list_table.delete(yql_client, oldest_date)

        crypta_table = InterestingCryptaTableDescription(
            table_path=os.path.join(YDB_GEO_ROOT, oldest_crypta),
            ydb_connection_params=ydb_connection_params
        )
        crypta_table.drop(yql_client)

        geo_table = YdbGeoLogsTableDescription(
            table_path=os.path.join(YDB_GEO_ROOT, oldest_geo),
            ydb_connection_params=ydb_connection_params
        )
        geo_table.drop(yql_client)

    return [task.make_done()]
