import shutil
import os
import json
from pathlib import Path
from datetime import datetime
from mlpipeline.entity import ModelPusherConfig, ModelPusherArtifact
from mlpipeline.logging.logger import get_logger
from mlpipeline.exception import ModelPusherException
import sys

logger = get_logger(__name__)


class ModelPusher:
    def __init__(self, config: ModelPusherConfig):
        self.config = config
    
    def push_model(self) -> ModelPusherArtifact:
        try:
            logger.info("Starting model pusher")
            
            os.makedirs(self.config.model_registry_path, exist_ok=True)
            os.makedirs(self.config.root_dir, exist_ok=True)
            
            model_source = Path(self.config.model_path)
            model_dest = Path(self.config.model_registry_path) / model_source.name
            
            if model_source.is_dir():
                if model_dest.exists():
                    shutil.rmtree(model_dest)
                shutil.copytree(model_source, model_dest)
            else:
                shutil.copy2(model_source, model_dest)
            
            logger.info(f"Model pushed to: {model_dest}")
            
            # Save push metadata
            metadata = {
                "pushed_at": datetime.now().isoformat(),
                "source_path": str(model_source),
                "destination_path": str(model_dest),
                "status": "success"
            }
            
            metadata_path = Path(self.config.root_dir) / "push_metadata.json"
            with open(metadata_path, 'w') as f:
                json.dump(metadata, f, indent=4)
            
            return ModelPusherArtifact(
                pushed_model_path=str(model_dest),
                is_pushed=True,
                message="Model pushed successfully"
            )
        except Exception as e:
            raise ModelPusherException(str(e), sys)