import os

from funcy import rcompose

from image_prediction.config import CONFIG
from image_prediction.default_objects import get_extractor_classifier, get_formatter, get_mlflow_model_loader
from image_prediction.locations import MLRUNS_DIR

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"


def load_pipeline(**kwargs):
    model_loader = get_mlflow_model_loader(MLRUNS_DIR)
    model_identifier = CONFIG.service.run_id

    pipeline = Pipeline(model_loader, model_identifier, progress_message="Processing document", **kwargs)

    return pipeline


class Pipeline:
    def __init__(self, model_loader, model_identifier, **kwargs):
        self.pipe = rcompose(get_extractor_classifier(model_loader, model_identifier, **kwargs), get_formatter())

    def __call__(self, pdf: bytes, page_range: range = None):
        yield from self.pipe(pdf, page_range=page_range)