from chandanops.core.logger import setup_logger
from chandanops.core.mlflow_manager import MLflowManager


class Pipeline:

    def __init__(self, experiment_name="chandanops"):
        self.logger = setup_logger()
        self.mlflow = MLflowManager(experiment_name)

    def run(self, train_fn, params):
        self.logger.info("Pipeline started")

        with self.mlflow.start_run():
            self.mlflow.log_params(params)

            model, metrics = train_fn(params)

            self.mlflow.log_metrics(metrics)
            self.mlflow.log_model(model)

        self.logger.info("Pipeline completed")
