from src.entity import config_entity
from src.entity import artifact_entity
from src.logger import logging
from src.exception import FertilizerException
from src.predictor import ModelResolver
from src.utils import load_object
from src.utils import save_object

from src.entity.config_entity import ModelPusherConfig

from src.entity.artifact_entity import DataTransformationArtifact
from src.entity.artifact_entity import ModelTrainerArtifact
from src.entity.artifact_entity import ModelPusherArtifact

import os 
import sys 

class ModelPusher:

    def __init__(
        self, 
        model_pusher_config: ModelPusherConfig,
        data_transformation_artifact: DataTransformationArtifact,
        model_trainer_artifact: ModelTrainerArtifact
    ):

        try:
            logging.info(f"\n\n{'>'*50} Model Pusher Initiated {'<'*50}\n")
            self.model_pusher_config = model_pusher_config
            self.data_transformation_artifact = data_transformation_artifact
            self.model_trainer_artifact = model_trainer_artifact
            self.model_resolver = ModelResolver(model_registry=self.model_pusher_config.saved_model_dir)
        
        except Exception as e:
            raise FertilizerException(e, sys)

    def initiate_model_pusher(self) -> ModelPusherArtifact:
        try:
            # load object
            logging.info(f"Loading transformer model and target encoder")
            transformer = load_object(file_path=self.data_transformation_artifact.transform_object_path)
            model = load_object(file_path=self.model_trainer_artifact.model_path)
            target_encoder = load_object(file_path=self.data_transformation_artifact.target_encoder_path)

            # model pusher dir 
            logging.info(f"Saving model into model pusher directory")
            save_object(file_path=self.model_pusher_config.pusher_transformer_path, obj=transformer)
            save_object(file_path=self.model_pusher_config.pusher_model_path, obj=model)
            save_object(file_path=self.model_pusher_config.pusher_target_encoder_path, obj=target_encoder)

            # saved model dir 
            logging.info(f"Saving model in saved model dir")

            transformer_path = self.model_resolver.get_latest_save_transformer_path()
            model_path = self.model_resolver.get_latest_save_model_path()
            target_encoder_path = self.model_resolver.get_latest_save_target_encoder_path()

            save_object(file_path=transformer_path, obj=transformer)
            save_object(file_path=model_path, obj=model)
            save_object(file_path=target_encoder_path, obj=target_encoder)

            model_pusher_artifact = ModelPusherArtifact(
                pusher_model_dir = self.model_pusher_config.pusher_model_dir, 
                saved_model_dir = self.model_pusher_config.saved_model_dir)
            
            logging.info(f"Model Pusher Artifacts Generated")

            return model_pusher_artifact
        
        except Exception as e:
            raise FertilizerException(e, sys)