import os

from datacloud.score_api.storage.cache.ttl_cache import TtlCache
from datacloud.score_api.storage.ydb.ydb_tables.score_tables.crypta_table import CryptaTable
from datacloud.score_api.storage.ydb.ydb_tables.config_tables.geo_path_table import GeoPathTable
from datacloud.score_api.storage.ydb.ydb_tables.geo_tables.geo_logs_table import GeoLogsTable


class YdbGeoStorage:
    def __init__(self, ydb_manager, database, geo_path_table_path, geo_root_path):
        self._ydb_manager = ydb_manager
        self._database = database

        self._geo_path_table_path = geo_path_table_path
        self._geo_path_table = None
        self._geo_path_cache = TtlCache(111)

        self._geo_root_path = geo_root_path

    @property
    def geo_path_table(self):
        if self._geo_path_table is None:
            self._geo_path_table = GeoPathTable(self._ydb_manager, self._database, self._geo_path_table_path)
        return self._geo_path_table

    def get_geo_paths_with_cache(self, internal_score_name):
        cache_key = (internal_score_name, )
        status, val = self._geo_path_cache.get(cache_key)
        if status == TtlCache.Status.OK:
            return val

        geo_paths_record = self.geo_path_table.get_one(GeoPathTable.Record(internal_score_name))
        if not geo_paths_record:
            return None

        val = geo_paths_record
        self._geo_path_cache.insert(cache_key, val)

        return val

    def get_crypta_path(self, internal_score_name):
        cache = self.get_geo_paths_with_cache(internal_score_name)
        if cache is None:
            return None

        return os.path.join(self._geo_root_path, cache.crypta_path)

    def get_geo_path(self, internal_score_name):
        cache = self.get_geo_paths_with_cache(internal_score_name)
        if cache is None:
            return None

        return os.path.join(self._geo_root_path, cache.geo_path)

    def get_geo_weights(self, internal_score_name):
        cache = self.get_geo_paths_with_cache(internal_score_name)
        if cache is None:
            return None

        return map(float, cache.weights.split())

    def get_hashed_cids_by_id_values(self, internal_score_name, id_values_list):
        crypta_path = self.get_crypta_path(internal_score_name)
        if crypta_path is None:
            return None

        crypta_table = CryptaTable(self._ydb_manager, self._database, crypta_path)
        return set(rec.hashed_cid for rec in crypta_table.get_multiple(id_values_list))

    def get_geo_logs_by_id_values(self, internal_score_name, id_values_list):
        geo_path = self.get_geo_path(internal_score_name)
        if geo_path is None:
            return None

        geo_logs_table = GeoLogsTable(self._ydb_manager, self._database, geo_path)
        hashed_cids = self.get_hashed_cids_by_id_values(internal_score_name, id_values_list)
        if not hashed_cids:
            return set()

        points = set()
        for rec in geo_logs_table.get_multiple(list(hashed_cids)):
            points.update(set(rec.points.split()))

        return points
