import pandas as pd

from typing import Tuple, Optional, Union, Dict
from pathlib import Path
from emissary_ml_toolkit.utils.logging import get_logger
from emissary_ml_toolkit.config.loader import ConfigLoader
from emissary_ml_toolkit.config.dataset import DatasetConfig, ProcessedOutputConfig, PreprocessingConfig
from emissary_ml_toolkit.preprocessing import transforms
from emissary_ml_toolkit.storage.manager import StorageManager
from emissary_ml_toolkit.webhooks.client import WebhookClient
from emissary_ml_toolkit.preprocessing.loaders import load_file

logger = get_logger(__name__)


class DatasetPipeline:
    """High-level pipeline for configuration-driven preprocessing

    This pipeline supports two distinct modes of operation:

    1. **Production Mode**: Load datasets from remote storage using configuration files.
       Intended for training jobs initiated from the platform with full storage and
       webhook integration.

    2. **Local Mode**: Load datasets from local file paths for experimentation.
       Intended for ML engineers testing preprocessing logic locally without
       platform integration

    Examples:
        Production mode (with configuration):
            >>> pipeline = DatasetPipeline.from_config("training_job.yaml")
            >>> train_df, test_df = pipeline.prepare_dataset_from_config()

        Local mode (with file paths)
            >>> pipeline = DatasetPipeline()
            >>> train_df, test_df = pipeline.prepare_dataset_from_local(
            ...     train_path="train.jsonl",
            ...     downsample=10,
            ... )
    """

    def __init__(self):
        self._storage_manager: Optional[StorageManager] = None
        self._webhook_client: Optional[WebhookClient] = None
        self._dataset_config: Optional[DatasetConfig] = None
        self._config_mode: bool = False
        self._config_path: Optional[Path] = None

    @classmethod
    def from_config(cls, config_path: Union[str, Path]) -> 'DatasetPipeline':
        """Create a pipeline instance configured for production mode

        Args:
            config_path: Path to the training job configuration YAML file

        Returns:
            DatasetPipeline instance initialized with storage and webhook clients

        Raises:
            FileNotFoundError: If config file does not exist
            ValueError: If config file is invalid
        """
        pipeline = cls()
        pipeline._initialize_from_config(config_path)
        return pipeline

    def prepare_dataset_from_config(
            self,
            config_path: Union[str, Path] = None,
    ) -> Tuple[pd.DataFrame, Optional[pd.DataFrame]]:
        """Prepare datasets using configuration file (production mode)

        This method loads datasets from remote storage, applies configured
        preprocessing, uploads results back to storage, and sends webhook
        notifications to the platform server.

        Args:
            config_path: Path to configuration file. If None, uses the path
                        provided during initialization via from_config()

        Returns:
            Tuple of (train_dataframe, test_dataframe). test_dataframe may be None
            if no test dataset is configured.

        Raises:
            ValueError: If no config_path provided and pipeline not initialized with one
            RuntimeError: If dataset loading or processing fails
        """
        if config_path is None and not self._config_mode:
            raise ValueError(
                "No configuration provided. Either pass config_path or use "
                "DatasetPipeline.from_config() to initialize the pipeline"
            )

        if config_path and not self._config_mode:
            self._initialize_from_config(config_path)

        logger.info("Starting dataset preparation from configuration")

        try:
            train_data, test_data, preprocessing_config = self._load_data_from_config()

            if preprocessing_config is None:
                logger.info("No preprocessing configured, returning raw datasets")
                return train_data, test_data

            train_dataset, test_dataset = self._preprocess_dataset(
                train_data=train_data,
                test_data=test_data,
                preprocessing_config=preprocessing_config,
            )

            logger.info("Dataset preparation completed successfully")
            return train_dataset, test_dataset
        except Exception as e:
            logger.error(f"Dataset preparation failed: {e}")
            raise RuntimeError(f"Failed to prepare dataset from config: {e}") from e

    def prepare_dataset_from_local(
            self,
            train_data_path: Union[str, Path],
            test_data_path: Union[str, Path] = None,
            downsample: Optional[int] = None,
            rebalance: Optional[Dict[str, Union[int, float]]] = None,
            split_ratio: Optional[int] = None,
    ):
        """Prepare datasets from local file paths (local/experimental mode)

        This method loads datasets from local files and applies preprocessing.
        No storage upload or webhook notifications are performed.

        Args:
            train_data_path: Path to training dataset file
            test_data_path: Optional path to test dataset file
            downsample: Percentage to downsample (0-100).
            rebalance: Dictionary mapping class labels to target ratios. (Only when task_type == classification)
            split_ratio: Percentage of data to use for test split (0-100).

        Returns:
            Tuple of (train_dataframe, test_dataframe). test_dataframe may be None

        Raises:
            FileNotFoundError: If specified file paths do not exist
            ValueError: If both downsample and rebalance are specified
            RuntimeError: If dataset loading or processing fails
        """
        logger.info(f"Starting dataset preparation from local files: {train_data_path}")

        try:
            train_data = load_file(train_data_path)
            test_data = load_file(test_data_path) if test_data_path else None

            preprocessing_config = PreprocessingConfig(
                downsample=downsample,
                rebalance=rebalance,
                split_ratio=split_ratio,
            ) if any([downsample, rebalance, split_ratio]) else None

            if preprocessing_config is None:
                logger.info("No preprocessing configured, returning raw datasets")
                return train_data, test_data

            train_dataset, test_dataset = self._transform_data(
                train_data=train_data,
                test_data=test_data,
                config=preprocessing_config,
            )

            logger.info("Dataset preparation completed successfully")
            return train_dataset, test_dataset

        except Exception as e:
            logger.error(f"Dataset preparation from files failed: {e}")
            raise RuntimeError(f"Failed to prepare dataset from files: {e}") from e

    def _initialize_from_config(self, config_path: Union[str, Path]) -> None:
        """Initialize storage manager and webhook client from configuration

        Args:
            config_path: Path to training job configuration file

        Raises:
            FileNotFoundError: If config file does not exist
            ValueError: If config file is malformed
        """
        config_path = Path(config_path)
        if not config_path.exists():
            raise FileNotFoundError(f"Configuration file not found: {config_path}")

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

        try:
            training_job_configs = ConfigLoader.from_yaml(config_path)
            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._dataset_config = training_job_configs.dataset_config
            self._config_mode = True
            self._config_path = config_path
            logger.info("DatasetPipeline initialized successfully in production mode")
        except Exception as e:
            logger.error(f"Failed to initialize DatasetPipeline from config: {e}")
            raise ValueError(f"Invalid configuration file: {e}") from e

    def _load_data_from_config(self) -> Tuple[pd.DataFrame, Optional[pd.DataFrame], Optional[PreprocessingConfig]]:
        train_dataset_config = self._dataset_config.train_dataset
        test_dataset_config = self._dataset_config.test_dataset

        train_data = self._load_file_from_storage(train_dataset_config)
        test_data = self._load_file_from_storage(test_dataset_config) if test_dataset_config else None

        preprocessing_config = train_dataset_config.preprocessing

        return train_data, test_data, preprocessing_config

    def _preprocess_dataset(
            self,
            train_data: pd.DataFrame,
            test_data: pd.DataFrame,
            preprocessing_config: PreprocessingConfig
    ):
        train_dataset, test_dataset = self._transform_data(
            train_data=train_data,
            test_data=test_data,
            config=preprocessing_config,
        )

        if preprocessing_config.output_config:
            self._upload_and_notify_processed_data(
                train_dataset=train_dataset,
                test_dataset=test_dataset,
                output_config=preprocessing_config.output_config
            )

        return train_dataset, test_dataset

    def _upload_and_notify_processed_data(
            self,
            train_dataset: pd.DataFrame,
            test_dataset: pd.DataFrame,
            output_config: ProcessedOutputConfig
    ):
        self._upload_datasets_to_storage(output_config=output_config, train_df=train_dataset, test_df=test_dataset)
        self._send_notifications(output_config=output_config, has_test=test_dataset is not None)

    def _transform_data(
            self,
            train_data: pd.DataFrame,
            test_data: pd.DataFrame,
            config: PreprocessingConfig
    ) -> Tuple[pd.DataFrame, Optional[pd.DataFrame]]:
        """Apply configured preprocessing (1. Split 2. Downsample or Rebalance)"""
        if config.downsample and config.rebalance:
            raise ValueError("Downsample and rebalance cannot be configured at the same time")

        if config.split_ratio and test_data is None:
            train_data, test_data = transforms.split(
                data=train_data,
                test_size=config.split_ratio / 100
            )

        if config.downsample:
            train_data = transforms.downsample(
                data=train_data,
                sample_ratio=config.downsample / 100
            )

        if config.rebalance:
            train_data = transforms.rebalance(train_data, config.rebalance)

        return train_data, test_data

    def _load_file_from_storage(self, dataset_config: DatasetConfig) -> pd.DataFrame:
        """Load dataset from configuration"""
        if self._storage_manager is None:
            raise RuntimeError(
                "Storage manager not initialized, Use from_config() to initialize "
                "the pipeline in production mode"
            )
        storage = self._storage_manager.get_storage(dataset_config.storage_ref)
        file_path = storage.download_file(dataset_config.storage_path)
        return load_file(file_path)

    def _upload_datasets_to_storage(
        self,
        output_config: ProcessedOutputConfig,
        train_df: pd.DataFrame,
        test_df: Optional[pd.DataFrame] = None,
    ) -> None:
        """Save processed datasets to configured storage"""
        storage = self._storage_manager.get_storage(output_config.storage_ref)

        train_path = f"{output_config.base_path}/{output_config.train_dataset}"
        storage.upload_dataframe(train_df, train_path)

        if test_df is not None:
            test_path = f"{output_config.base_path}/{output_config.test_dataset}"
            storage.upload_dataframe(test_df, test_path)

    def _send_notifications(self, output_config: ProcessedOutputConfig, has_test: bool):
        """Send webhook notifications for uploaded datasets"""
        train_dataset_id = output_config.train_dataset.split('/')[0]
        self._webhook_client.send_dataset_upload(train_dataset_id)

        if has_test:
            test_dataset_id = output_config.test_dataset.split('/')[0]
            self._webhook_client.send_dataset_upload(test_dataset_id)

        return
