import logging
import os
import re
import torch
import threading
import json
import numpy as np
from datetime import datetime
from fitxf import TensorUtils, FitXformInterface
from fitxf.math.algo.encoding.Base64 import Base64
from fitxf.math.datasource.vecdb.model.ModelEncoderInterface import ModelEncoderInterface
from fitxf.math.datasource.vecdb.model.ModelDbInterface import ModelDbInterface
#from fitxf.math.datasource.vecdb.model.ModelDb import ModelDb
from fitxf.math.datasource.vecdb.metadata.MetadataInterface import MetadataInterface
from fitxf.math.utils.Lock import Lock


class ModelInterface:

    OLD_DATETIME = datetime(year=1999, month=9, day=9)

    def __init__(
            self,
            user_id: str,
            llm_model: ModelEncoderInterface,
            col_content: str,
            col_label_user: str,
            col_label_std: str,
            col_embedding: str,
            numpy_to_b64_for_db: bool,
            vec_db_metadata: MetadataInterface,
            fit_xform_model: FitXformInterface,
            cache_tensor_to_file: bool,
            file_temp_dir: str,
            # allowed values: "np", "torch"
            return_tensors: str,
            enable_bg_thread_for_training: bool = True,
            logger = None,
    ):
        self.user_id = user_id
        self.llm_model = llm_model
        self.col_content = col_content
        self.col_label_user = col_label_user
        self.col_label_standardized = col_label_std
        self.col_embedding = col_embedding
        self.numpy_to_b64_for_db = numpy_to_b64_for_db
        self.vec_db_metadata = vec_db_metadata
        self.fit_xform_model = fit_xform_model
        self.cache_tensor_to_file = cache_tensor_to_file
        self.file_temp_dir = file_temp_dir
        self.return_tensors = return_tensors
        self.enable_bg_thread_for_training = enable_bg_thread_for_training
        self.logger = logger if logger is not None else logging.getLogger()

        self.base64_encoder = Base64(logger=self.logger)
        self.tensor_utils = TensorUtils(logger=self.logger)
        self.llm_model_path = self.llm_model.get_model_path()

        self.model_db = self.__get_model_db(ModelClass=ModelDb)
        self.logger.info('DB params "' + str(self.user_id) +'": ' + str(self.model_db.get_db_params().get_db_info()))

        # Lock only by threads in the same worker
        self.mutex_name_model = 'model'
        self.mutex_name_underlying_db = 'db'
        self.lock_mutexes = Lock(
            # User may also provide their own storage (in the usual case)
            mutex_names = [self.mutex_name_model, self.mutex_name_underlying_db],
            logger = self.logger,
        )

        assert os.path.isdir(self.file_temp_dir), 'Not a directory for tmp caching "' + str(self.file_temp_dir) + '"'
        # Don't use characters like "." because during cleanup cant find
        self.file_name_prefix_for_match_pattern = 'VecDbModelStandard_' + \
                                str(datetime.now().strftime("%Y%m%d_%H%M%S")) + '_' \
                                + str(np.random.rand())[2:]
        self.file_name_prefix = self.file_temp_dir + '/' + self.file_name_prefix_for_match_pattern

        self.last_sync_time_with_underlying_db = self.OLD_DATETIME
        self.logger.info(
            'LLM model path "' + str(self.llm_model_path)
            + '", encode numpy to base 64 in DB = ' + str(self.numpy_to_b64_for_db)
        )
        
        self.bg_thread = threading.Thread(target=self.run_bg_thread)
        self.bg_thread_sleep_secs = float(os.environ.get("VECDB_BG_SLEEP_SECS", "10."))
        self.clear_memory_secs_inactive = float(os.environ.get("VECDB_CLEAR_MEMORY_SECS_INACTIVE", "30."))
        self.signal_stop_bg_thread = False
        self.logger.info(
            'Set bg thread sleep to ' + str(self.bg_thread_sleep_secs)
            + 's, clear memory inactive time to ' + str(self.clear_memory_secs_inactive) + 's.'
        )
        return

    def __get_model_db(
            self,
            ModelClass: type(ModelDbInterface),
    ) -> ModelDbInterface:
        return ModelClass(
            tablename = self.user_id,
            col_content = self.col_content,
            col_label_user = self.col_label_user,
            col_label_standardized = self.col_label_standardized,
            col_embedding = self.col_embedding,
            max_records = 999999,
            logger = self.logger,
        )

    def check_model_consistency_with_prev(
            self,
    ):
        # Verify that model name has not changed from previous
        try:
            metadata_model_json = json.loads(self.vec_db_metadata.get_metadata_model_last_update())
            model_prev = metadata_model_json['llm_model']
        except Exception as ex:
            self.logger.error('Error getting metadata model: ' + str(ex))
            model_prev = None
        self.logger.info('Previous model from metadata "' + str(model_prev) + '"')

        # csv can return float type "nan" instead of None
        if model_prev is not None:
            name_prev = self.get_model_name_from_path(model_path=model_prev)
            name_cur = self.get_model_name_from_path(model_path=self.llm_model_path)
            assert name_prev == name_cur, \
                'Previous model different "' + str(model_prev) + '" (type "' + str(type(model_prev)) + \
                '") from new "' + str(self.llm_model_path) + '" (type "' + str(self.llm_model_path) \
                + '")'
            self.logger.info(
                'Model consistency check ok, cur/prev name same "' + str(name_cur)
                + '", full path "' + str(self.llm_model_path) + '"'
            )
        else:
            self.logger.warning('Model from metadata returned None')

    def cleanup(self):
        try:
            files_in_tmp_dir = os.listdir(path=self.file_temp_dir)
            self.logger.info('Files in temp dir "' + str(self.file_temp_dir) + '": ' + str(files_in_tmp_dir))
            for f in files_in_tmp_dir:
                if re.match(pattern=str(self.file_name_prefix_for_match_pattern)+".*", string=f):
                    self.logger.info(
                        'Cleanup deleting file "' + str(f) + '" using prefix match pattern "'
                        + str(self.file_name_prefix_for_match_pattern) + '"'
                    )
                    os.remove(path=f)
                else:
                    self.logger.debug(
                        'Cleanup ignoring file "' + str(f) + '" using prefix match pattern "'
                        + str(self.file_name_prefix_for_match_pattern) + '"'
                    )
        except Exception as ex:
            self.logger.error('Error cleaning up temp dir "' + str(self.file_temp_dir) + '": ' + str(ex))
        return

    def stop_threads(self):
        if not self.signal_stop_bg_thread:
            self.signal_stop_bg_thread = True
            self.bg_thread.join()
            self.cleanup()
            self.logger.info('All threads successfully stopped')
        return

    def run_bg_thread(
            self,
    ):
        raise Exception('Must be implemented by derived class')

    def calc_embedding(
            self,
            content_list,
    ):
        assert self.llm_model is not None, 'LLM model is None'
        text_encode = self.llm_model.encode(
                content_list = content_list,
                return_tensors = self.return_tensors,
            )
        return text_encode

    def get_model_name_from_path(self, model_path):
        parts = [name for name in str(model_path).strip().split(sep=os.sep) if len(name) > 0]
        self.logger.info('Model path separated as ' + str(parts))
        return parts[-1]

    def update_metadata_model_updated(
            self,
            model_save_b64json_string = None,
    ):
        # CSV still got problem with json string
        self.vec_db_metadata.update_metadata_model_updated(
            llm_path = self.llm_model_path,
            model_save_b64json_string = model_save_b64json_string,
        )
        return

    def update_metadata_db_data_updated(self):
        self.vec_db_metadata.update_metadata_db_raw_data_updated()
        return

    def init_data_model(
            self,
            max_tries = 1,
            background = False,
    ):
        raise Exception('Must be implemented by derived class')

    def load_data_model(
            self,
            max_tries = 1,
            background = False,
    ):
        raise Exception('Must be implemented by derived class')

    def reset_data_model__(self):
        self.map_lbl_to_idx = {}
        self.map_idx_to_lbl = {}
        self.text_labels_standardized = np.array([])
        self.data_records = []
        return

    def update_model(
            self,
            force_update = False,
    ):
        raise Exception('Must be implemented by derived class')

    def set_high_urgency_to_sync_with_underlying_db(self):
        self.last_sync_time_with_underlying_db = self.OLD_DATETIME

    def get_data_length(self):
        return len(self.text_labels_standardized)

    def atomic_delete_add(
            self,
            delete_key: str,
            # list of dicts
            records: list[dict],
    ):
        raise Exception('Must be implemented by derived class')

    def add(
            self,
            # list of dicts
            records: list,
    ):
        raise Exception('Must be implemented by derived class')

    def delete(
            self,
            match_phrases,
    ):
        raise Exception('Must be implemented by derived class')

    def predict(
            self,
            text_list_or_embeddings,
            top_k = 5,
            # Instead of just returning the user labels, return full record. Applicable to some models only
            return_full_record = False,
    ):
        raise Exception('Must be implemented by derived class')

    def is_updating_model(self):
        return self.lock_mutexes.is_locked(mutex=self.mutex_name_model)

    ############################################################################################################
    # Underlying DB
    ############################################################################################################

    def sync__db(
            self,
            max_tries = 1,
            background = False,
    ):
        raise Exception('Must be implemented by derived class')

    # Heuristic measure from 0 to 1, based on secs since last load time (proportional),
    # training data size (inversely proportional), etc
    def is_need_sync_db(
            self,
    ):
        # secs_since_last_load = Profiling().get_time_dif_secs(
        #     start = self.last_sync_time_with_underlying_db,
        #     stop = last_updated_time_underlying_db
        # )
        db_last_update_time = self.vec_db_metadata.get_metadata_db_data_last_update()
        self.logger.debug(
            'Last DB update time "' + str(db_last_update_time)
            + '", last sync time ' + str(self.last_sync_time_with_underlying_db)
        )
        if db_last_update_time is None:
            need_update = True
        else:
            need_update = db_last_update_time > self.last_sync_time_with_underlying_db

        if need_update:
            self.logger.info(
                'Need update. Last DB update time "' + str(db_last_update_time)
                + '", last sync time ' + str(self.last_sync_time_with_underlying_db)
            )
        return need_update

    def is_loading_model_from_underlying_db(self):
        return self.lock_mutexes.is_locked(mutex=self.mutex_name_underlying_db)

    def get_text_encoding_from_db_records(
            self,
            db_records,
    ) -> np.ndarray:
        # at the same time, remove embedding record from db_records
        text_encoded = [r.pop(self.col_embedding) for r in db_records]
        if self.numpy_to_b64_for_db:
            text_encoded = [
                # 3 steps: Encode from base 64 string to base 64 bytes, then decode the base 64 to actual bytes,
                # then finally convert to numpy from base 64 bytes
                self.base64_encoder.decode_base64_string_to_numpy_array(s64=s64, data_type='float64')
                for s64 in text_encoded
            ]

        if (not self.numpy_to_b64_for_db) and (self.model_db.get_db_params().db_type == 'csv'):
            text_encoded_float = []
            for row in text_encoded:
                float_array = self.model_db.convert_csv_string_array_to_float_array(
                    string_array = row,
                    custom_chars_remove = [','],
                )
                text_encoded_float.append(float_array)
            text_encoded = text_encoded_float
        return np.array(text_encoded)

    def convert_to_embeddings_if_necessary(
            self,
            text_list_or_embeddings,
    ):
        txt_lm = None
        if type(text_list_or_embeddings) is np.ndarray:
            if text_list_or_embeddings.dtype in [
                'float64', 'float32', 'float', 'int64', 'int32', 'int',
                torch.float64, torch.float32, torch.float, torch.int64, torch.int32, torch.int,
            ]:
                txt_lm = text_list_or_embeddings
            else:
                self.logger.info(
                    'Passed in data type "' + str(text_list_or_embeddings.dtype) + '", assuming is text type'
                )

        if txt_lm is None:
            self.logger.info('Passed in data is likely not embedding but text list, calculating embedding..')
            txt_lm = self.calc_embedding(
                content_list = text_list_or_embeddings,
            )
        return txt_lm

    def __update_label_mapping(
            self,
            # when new data added, we might need to create new label mapping
            labels_from_new_records,
    ):
        # Problem with mapping to ordered integer is that in multithreaded environment like Rest Server,
        # race condition can cause 2 labels to generate same index, if there is no locking within programs
        # in different workers or servers.
        # But if we generate random numbers, then the probability of that happening is almost 0.
        for lb in labels_from_new_records:
            if lb not in self.map_lbl_to_idx:
                while True:
                    int_tmp = np.random.randint(low=0, high=2**31)
                    if int_tmp in self.map_idx_to_lbl.keys():
                        continue
                    self.map_lbl_to_idx[lb] = int_tmp
                    break
        self.map_idx_to_lbl = {i: lbl for lbl, i in self.map_lbl_to_idx.items()}
        return

    def update_label_maps_from_new_recs__(
            self,
            records,
            text_encoding_tensor,
    ):
        labels = [row[self.col_label_user] for row in records]
        self.__update_label_mapping(labels_from_new_records=labels)

        # Save records in append mode
        records_with_embedding = []
        for rec, txt_enc in zip(records, text_encoding_tensor):
            lbl_user = rec[self.col_label_user]
            # add standardized label to existing object in passed in records
            rec[self.col_label_standardized] = self.map_lbl_to_idx.get(lbl_user, -1)

            # Create a new record to be appended, don't add embedding to passed in records
            rec_copy = {k:v for k,v in rec.items()}
            rec_copy[self.col_embedding] = txt_enc
            records_with_embedding.append(rec_copy)

        return records_with_embedding

    def add_records_to_underlying_db__(
            self,
            records_with_embedding_and_labelstd,
    ):
        if self.numpy_to_b64_for_db:
            recs_for_storage_base64 = []
            for rec in records_with_embedding_and_labelstd:
                r = rec.copy()
                r[self.col_embedding] = self.base64_encoder.encode_numpy_array_to_base64_string(
                    x = r[self.col_embedding],
                    data_type = 'float64',
                )
                recs_for_storage_base64.append(r)
            self.logger.info(
                'Successfully encoded all embedding to base 64 string for storage.'
            )
        else:
            recs_for_storage_base64 = records_with_embedding_and_labelstd

        self.model_db.insert(
            records = recs_for_storage_base64,
            tablename = self.user_id,
        )
        self.logger.info(
            'Train update underlying DB saved ' + str(len(records_with_embedding_and_labelstd)) + ' records: '
            + str([{k: v for k, v in d.items() if k != self.col_embedding} for d in recs_for_storage_base64])
        )
        return

    def delete_records_from_underlying_db__(
            self,
            match_phrases,
    ):
        total_deleted = 0
        # Get all encoded texts, standardized labels, text encoding from memory cache
        for mp in match_phrases:
            self.__delete_single_record_from_underlying_db(
                match_phrase = mp,
            )
        return total_deleted

    def __delete_single_record_from_underlying_db(
            self,
            match_phrase,
    ):
        total_deleted = 0
        assert len(match_phrase) == 1, 'Not supported delete with more than 1 key phrase ' + str(match_phrase)
        self.logger.info('Trying to delete records using match phrase: ' + str(match_phrase))
        res_del = self.model_db.delete(
            match_phrase = match_phrase,
            tablename = self.user_id,
        )
        self.logger.info('Deleted record with match phrase ' + str(match_phrase) + ', result ' + str(res_del))
        return 1


if __name__ == '__main__':
    exit(0)
