import importlib

from pandas import DataFrame

import mindsdb.interfaces.storage.db as db
from mindsdb.interfaces.storage.model_fs import ModelStorage, HandlerStorage
from mindsdb.integrations.libs.ml_handler_process.handlers_cacher import handlers_cacher
from mindsdb.utilities.functions import mark_process


@mark_process(name="learn")
def predict_process(
    integration_id: int,
    predictor_record: db.Predictor,
    args: dict,
    module_path: str,
    ml_engine_name: str,
    dataframe: DataFrame,
) -> DataFrame:
    module = importlib.import_module(module_path)

    if predictor_record.id not in handlers_cacher:
        handlerStorage = HandlerStorage(integration_id)
        modelStorage = ModelStorage(predictor_record.id)
        ml_handler = module.Handler(
            engine_storage=handlerStorage,
            model_storage=modelStorage,
        )
        handlers_cacher[predictor_record.id] = ml_handler
    else:
        ml_handler = handlers_cacher[predictor_record.id]

    predictions = ml_handler.predict(dataframe, args)
    ml_handler.close()
    return predictions
