from typing import Union, Optional, Tuple
from pathlib import Path
from transformers import Trainer
from emissary_ml_toolkit.utils.logging import get_logger
from emissary_ml_toolkit.config.loader import ConfigLoader
from emissary_ml_toolkit.config.schemas import TrainConfig, BaseModelConfig
from emissary_ml_toolkit.storage.manager import StorageManager
from emissary_ml_toolkit.webhooks.client import WebhookClient
from emissary_ml_toolkit.train.callbacks import EmissaryTrainerCallback
from emissary_ml_toolkit.train.model_downloader import ModelDownloader


logger = get_logger(__name__)


class TrainPipeline:
    """High-level pipeline for configuration-driven training

    This pipeline supports two distinct modes of operation:

    1. **Production Mode**: Support add_callback functioning which will adds callback to the
       given trainer which automatically handles artifacts upload / webhook notifications.

    2. **Local Mode**: Will have no interaction with remote storage and platform.
       Fetch train configurations from the yaml file.

    Examples:
        Production mode:
            >>> pipeline = TrainPipeline.from_config("training_job.yaml")
            >>> trainer = pipeline.add_callback(trainer)
            >>> trainer.train()
        Local mode:
            >>> pipeline = TrainPipeline.from_config("training_job.yaml", mode="local")
            >>> train_config = pipeline.get_train_config()
            >>> lora_config = pipeline.get_lora_config()
            >>> hyper_parameters = pipeline.get_hyper_parameters()
    """
    def __init__(self, mode: str):
        self.trainer = None
        self._storage_manager = None
        self._webhook_client = None
        self.train_config = None
        self.base_model_config = None
        self._model_downloader = ModelDownloader()
        self._storage = None
        self._mode = mode

    @classmethod
    def from_config(
            cls,
            config_path: Union[str, Path],
            mode: Optional[str] = "production",
    ) -> "TrainPipeline":
        pipeline = cls(mode=mode)
        pipeline._initialize_from_config(config_path)
        return pipeline

    def _initialize_from_config(self, config_path: Union[str, Path]) -> None:
        config_path = Path(config_path)
        if not config_path.exists():
            raise FileNotFoundError(f"Configuration file not found: {config_path}")

        logger.info(f"Initializing train pipeline from configuration: {config_path}")

        try:
            training_job_configs = ConfigLoader.from_yaml(config_path)
            self.train_config = training_job_configs.train_config
            self.base_model_config = training_job_configs.base_model_config
            if self._mode == "production":
                self._storage_manager = StorageManager(training_job_configs.storage_config)
                self._webhook_client = WebhookClient(
                    training_job_configs.project_id,
                    training_job_configs.training_job_id,
                )
                self._model_downloader = ModelDownloader(self._storage_manager)
                self._storage = self._storage_manager.get_storage(self.train_config.training_output.storage_ref)

            logger.info(f"TrainPipeline (mode={self._mode}) initialized successfully")

        except Exception as e:
            logger.error(f"Failed to initialize TrainPipeline from config: {e}")
            raise ValueError(f"Invalid configuration file: {e}") from e

    def get_train_config(self) -> TrainConfig:
        return self.train_config

    def get_base_model_config(self) -> BaseModelConfig:
        return self.base_model_config

    def get_lora_config(self) -> dict:
        return self.train_config.lora_config.model_dump()

    def get_hyper_parameters(self) -> dict:
        return self.train_config.hyper_parameters

    def download_model(self, local_dir: Union[str, Path] = "./base_model") -> Tuple[str, Optional[str]]:
        return self._model_downloader.download(self.base_model_config, local_dir)

    def add_callback(self, trainer: Trainer) -> Trainer:
        if self._mode != "production":
            raise ValueError(f"add_callback does not support in {self._mode} mode ")
        logger.info("Attaching essential callbacks to the trainer")
        emissary_trainer_callback = EmissaryTrainerCallback(
            storage=self._storage,
            webhook_client=self._webhook_client,
            train_config=self.train_config,
        )

        if not hasattr(trainer, 'add_callback'):
            raise AttributeError(
                f"Trainer object of type {type(trainer).__name__} does not have 'add_callback' method"
            )

        trainer.add_callback(emissary_trainer_callback)
        return trainer

    def send_success_notification(self):
        if self._mode != "production":
            return
        train_log_storage_path = f"{self.train_config.training_output.base_path}/train.log"
        self._storage.upload_file("train.log", train_log_storage_path)
        self._webhook_client.send_training_status("Success", True)

    def send_failure_notification(self):
        if self._mode != "production":
            return
        train_log_storage_path = f"{self.train_config.training_output.base_path}/train.log"
        self._storage.upload_file("train.log", train_log_storage_path)
        self._webhook_client.send_training_status("Failed", True)

