# -*- coding: utf-8 -*-
import os
import ydb
from concurrent.futures import TimeoutError
from datacloud.dev_utils.logging.logger import get_basic_logger
import datacloud.score_api_sample_data

from datacloud.score_api.storage.ydb import ydb_tables
from datacloud.score_api.storage.ydb import utils as ydb_utils

logger = get_basic_logger(__name__)


FillDataQuery = """PRAGMA TablePathPrefix("{}");

-- Config tables
DECLARE $partnerScoresData AS "List<Struct<
    partner_id: String,
    partner_score_name: String,
    internal_score_name: String,
    is_active: Bool>>";

DECLARE $partnerTokensData AS "List<Struct<
    partner_id: String,
    token: String>>";

DECLARE $scoreToPathData AS "List<Struct<
    internal_score_name: String,
    score_path: String>>";

DECLARE $cryptaTableData AS "List<Struct<
    hashed_id: Uint64,
    hashed_cid: Uint64>>";

DECLARE $scoreTableData AS "List<Struct<
    hashed_cid: Uint64,
    score: Double>>";

REPLACE INTO [datacloud-score-api/config/partner_scores]
SELECT
    partner_id,
    partner_score_name,
    internal_score_name,
    is_active
FROM AS_TABLE($partnerScoresData);

REPLACE INTO [datacloud-score-api/config/partner_tokens]
SELECT
    partner_id,
    token
FROM AS_TABLE($partnerTokensData);

REPLACE INTO [datacloud-score-api/config/score_path]
SELECT
    internal_score_name,
    score_path
FROM AS_TABLE($scoreToPathData);
-- End config tables


REPLACE INTO [datacloud-score-api/scores/test-score-1/23-10-2018/crypta]
SELECT
    hashed_id,
    hashed_cid
FROM AS_TABLE($cryptaTableData);

REPLACE INTO [datacloud-score-api/scores/test-score-1/23-10-2018/score]
SELECT
    hashed_cid,
    score
FROM AS_TABLE($scoreTableData);
"""


def fill_tables_with_data(session, path):
    global FillDataQuery

    prepared_query = session.prepare(FillDataQuery.format(path))
    session.transaction(ydb.SerializableReadWrite()).execute(
        prepared_query, {
            '$partnerScoresData': datacloud.score_api_sample_data.get_partner_scores_data(),
            '$partnerTokensData': datacloud.score_api_sample_data.get_partner_tokens_data(),
            '$scoreToPathData': datacloud.score_api_sample_data.get_score_path_data(),
            '$cryptaTableData': datacloud.score_api_sample_data.get_crypta_data(),
            '$scoreTableData': datacloud.score_api_sample_data.get_score_data()
        },
        commit_tx=True,
    )


def select_simple(session, path):
    # new transaction in serializable read write mode
    # if query successfully completed you will get result sets.
    # otherwise exception will be raised
    result_sets = session.transaction(ydb.SerializableReadWrite()).execute(
        """
        PRAGMA TablePathPrefix("{}");
        SELECT series_id, title, DateTime::ToDate(DateTime::FromDays(release_date)) AS release_date
        FROM series
        WHERE series_id = 1;
        """.format(path),
        commit_tx=True,
    )
    print("\n> select_simple_transaction:")
    for row in result_sets[0].rows:
        print("series, id: ", row.series_id, ", title: ", row.title, ", release date: ", row.release_date)

    return result_sets[0]


def upsert_simple(session, path):
    session.transaction().execute(
        """
        PRAGMA TablePathPrefix("{}");
        UPSERT INTO episodes (series_id, season_id, episode_id, title) VALUES
            (2, 6, 1, "TBD");
        """.format(path),
        commit_tx=True,
    )


def select_prepared(session, path, series_id, season_id, episode_id):
    query = """
    PRAGMA TablePathPrefix("{}");

    DECLARE $seriesId AS Uint64;
    DECLARE $seasonId AS Uint64;
    DECLARE $episodeId AS Uint64;

    SELECT title, DateTime::ToDate(DateTime::FromDays(air_date)) as air_date
    FROM episodes
    WHERE series_id = $seriesId AND season_id = $seasonId AND episode_id = $episodeId;
    """.format(path)

    prepared_query = session.prepare(query)
    result_sets = session.transaction(ydb.SerializableReadWrite()).execute(
        prepared_query, {
            '$seriesId': series_id,
            '$seasonId': season_id,
            '$episodeId': episode_id,
        },
        commit_tx=True
    )
    print("\n> select_prepared_transaction:")
    for row in result_sets[0].rows:
        print("episode title:", row.title, ", air date:", row.air_date)

    return result_sets[0]


# Show usage of explicit Begin/Commit transaction control calls.
# In most cases it's better to use transaction control settings in session.transaction
# calls instead to avoid additional hops to YDB cluster and allow more efficient
# execution of queries.
def explicit_tcl(session, path, series_id, season_id, episode_id):
    query = """
    PRAGMA TablePathPrefix("{}");

    DECLARE $seriesId AS Uint64;
    DECLARE $seasonId AS Uint64;
    DECLARE $episodeId AS Uint64;

    UPDATE episodes
    SET air_date = DateTime::ToDays(DateTime::TimestampFromString("2018-09-11T15:15:59.373006Z"))
    WHERE series_id = $seriesId AND season_id = $seasonId AND episode_id = $episodeId;
    """.format(path)
    prepared_query = session.prepare(query)

    print('Query is:\n{}'.format(prepared_query))

    # # Get newly created transaction id
    # tx = session.transaction(ydb.SerializableReadWrite()).begin()

    # # Execute data query.
    # # Transaction control settings continues active transaction (tx)
    # tx.execute(
    #     prepared_query, {
    #         '$seriesId': series_id,
    #         '$seasonId': season_id,
    #         '$episodeId': episode_id
    #     }
    # )

    # print("\n> explicit TCL call")

    # # Commit active transaction(tx)
    # tx.commit()


def create_datacloud.score_api_config_tables(session, driver, database_path, root_path):
    config_folder = os.path.join(root_path, 'config')
    ydb_utils.create_folder(driver, database_path, config_folder)
    partner_scores_table = os.path.join(database_path, config_folder, 'partner_scores')
    score_to_path_table = os.path.join(database_path, config_folder, 'score_path')
    partner_tokens_table = os.path.join(database_path, config_folder, 'partner_tokens')

    session.create_table(
        partner_scores_table,
        ydb.TableDescription()
        .with_column(ydb.Column('partner_id', ydb.OptionalType(ydb.DataType.String)))
        .with_column(ydb.Column('partner_score_name', ydb.OptionalType(ydb.DataType.String)))
        .with_column(ydb.Column('internal_score_name', ydb.OptionalType(ydb.DataType.String)))
        .with_column(ydb.Column('is_active', ydb.OptionalType(ydb.DataType.Bool)))
        .with_primary_keys('partner_id', 'partner_score_name')
    )
    session.create_table(
        score_to_path_table,
        ydb.TableDescription()
        .with_column(ydb.Column('internal_score_name', ydb.OptionalType(ydb.DataType.String)))
        .with_column(ydb.Column('score_path', ydb.OptionalType(ydb.DataType.String)))
        .with_primary_key('internal_score_name')
    )
    session.create_table(
        partner_tokens_table,
        ydb.TableDescription()
        .with_column(ydb.Column('partner_id', ydb.OptionalType(ydb.DataType.String)))
        .with_column(ydb.Column('token', ydb.OptionalType(ydb.DataType.String)))
        .with_primary_key('partner_id')
    )


def create_datacloud.score_api_config_tables_v2(session, driver, database_path, root_path):
    config_folder = os.path.join(root_path, 'config')
    ydb_utils.create_folder(driver, database_path, config_folder)
    partner_scores_table_path = os.path.join(database_path, config_folder, 'partner_scores')

    partner_scores_table = ydb_tables.PartnerScoresTable(database_path, partner_scores_table_path)
    partner_scores_table.create(session)


def create_score_tables(session, driver, database_path, root_path, score_name, date_str):
    scores_folder = os.path.join(root_path, 'scores')
    new_score_folder = os.path.join(scores_folder, score_name, date_str)
    ydb_utils.create_folder(driver, database_path, new_score_folder)

    crypta_table = os.path.join(database_path, new_score_folder, 'crypta')
    score_table = os.path.join(database_path, new_score_folder, 'score')
    session.create_table(
        crypta_table,
        ydb.TableDescription()
        .with_column(ydb.Column('hashed_id', ydb.OptionalType(ydb.DataType.Uint64)))
        .with_column(ydb.Column('hashed_cid', ydb.OptionalType(ydb.DataType.Uint64)))
        .with_primary_key('hashed_id')
    )
    session.create_table(
        score_table,
        ydb.TableDescription()
        .with_column(ydb.Column('hashed_cid', ydb.OptionalType(ydb.DataType.Uint64)))
        .with_column(ydb.Column('score', ydb.OptionalType(ydb.DataType.Double)))
        .with_primary_key('hashed_cid')
    )


def create_datacloud.score_api_tables(session, driver, database, root_path):
    create_datacloud.score_api_config_tables(session, driver, database, root_path)
    create_score_tables(session, driver, database, root_path, 'test-score-1', '23-10-2018')
    score_tables_list = [
        ('partner_a', 'score_a'),
        ('partner_a', 'score_b'),
        ('partner_a', 'score_c'),
        ('partner_b', 'score_a'),
        ('partner_b', 'score_b'),
        ('partner_b', 'score_c'),
        ('partner_c', 'score_a'),
        ('partner_c', 'score_b'),
        ('partner_c', 'score_c'),
    ]
    for partner, score in score_tables_list:
        create_score_tables(session, driver, database, root_path, os.path.join(partner, score), '23-10-2018')


# endpoint ydb-ru.yandex.net:2135


def run(endpoint, database, path, auth_token):
    connection_params = ydb.ConnectionParams(endpoint, database=database, auth_token=auth_token)
    try:
        driver = ydb.Driver(connection_params)
        driver.wait(timeout=5)
    except TimeoutError:
        raise RuntimeError("Connect failed to YDB")
    session = driver.table_client.session().create()
    ydb_utils.ensure_path_exists(driver, database, path)

    # root_path = 'datacloud-score-api'
    # +
    # create_datacloud.score_api_tables(session, driver, database, root_path)
    # ydb_utils.describe_table(session, database, os.path.join(root_path, 'config/partner_scores'))
    # fill_tables_with_data(session, database)

    root_path = 'datacloud-score-api-v2'
    create_datacloud.score_api_config_tables_v2(session, driver, database, root_path)
    
    # select_simple(session, database)
    # upsert_simple(session, database)
    # select_prepared(session, database, 2, 3, 7)
    # select_prepared(session, database, 2, 3, 8)
    # explicit_tcl(session, database, 2, 6, 1)
    # select_prepared(session, database, 2, 6, 1)


if __name__ == '__main__':
    endpoint = 'ydb-ru.yandex.net:2135'
    database = '/ru/home/re9ulusv/mydb'
    path = ''

    run(
        endpoint,
        database,
        path,
        os.environ['YDB_TOKEN']
    )
