"""
MLFlow Repository

This module contains the MLFlowRepository class, which is responsible for
handling model training artifacts and MLFlow operations for the Model Manager system.

It includes methods for generating training reports, managing artifacts,
and logging model runs to MLFlow.

"""

import os
import shutil
from datetime import datetime
from os import makedirs, path

import numpy as np
import pandas as pd
from sientia_do.observability.logger import Logger

from model_manager.sientia.model_serving import ModelServing  # type: ignore[import-untyped]
from model_manager.sientia.reports import Reports  # type: ignore[import-untyped]
from model_manager.utils.models.train_model_result import TrainModelResult


class ModelRepository:
    def __init__(self, url, username, password, logger: Logger):
        self.model_serving = ModelServing(tracking_uri=url, username=username, password=password)

        self.logger = logger

    def save_model(self, train_result: TrainModelResult) -> TrainModelResult:
        """
        Save a trained ML model and its artifacts to MLflow.

        This activity orchestrates the complete model saving pipeline:
        1. Generates the next run name for the experiment
        2. Creates and organizes artifacts (reports, data files)
        3. Logs model, parameters, metrics, and artifacts to MLflow

        Args:
            input_data: Configuration for model saving operation
                Required keys:
                    - metadata (dict): Workflow execution metadata
                    - train_result (TrainModelResult): Training result with model and metrics

        Returns:
            TrainModelResult: Updated training result with run_name and artifacts

        Raises:
            Exception: If model saving fails (after sending notification)
        """
        experiment_name = train_result.params.experiment_name
        self.logger.info(f'Starting model save for experiment: {experiment_name}')

        # Step 1: Generate next run name
        self.logger.info('Generating run name')
        train_result.run_name = self._get_next_run_name(experiment_name)
        self.logger.info(f'Generated run name: {train_result.run_name}')

        # Step 2: Generate artifacts (reports, CSV files)
        self.logger.info('Generating artifacts')
        train_result = self._generate_artifacts(train_result)
        self.logger.info('Artifacts generated successfully')

        # Step 3: Save run to MLflow
        self.logger.info('Saving run to MLflow')
        self._save_run(train_result)

        self.logger.info(
            f'Model saved successfully - Run: {train_result.run_name}, '
            f'Experiment: {experiment_name}'
        )

        return train_result

    def cleanup_run_directory(self, run_dir: str) -> None:
        """
        Clean up temporary run directory after model training.

        This activity deletes the temporary directory created during model training
        and artifact generation. It implements idempotent cleanup to handle cases
        where the directory may have already been deleted.

        Args:
            run_dir (str): Path to the run directory to delete
        """
        if not run_dir:
            self.logger.info('No run directory specified, skipping cleanup')
            return

        self.logger.info(f'Cleaning up run directory: {run_dir}')

        if os.path.exists(run_dir):
            shutil.rmtree(run_dir)
            self.logger.info(f'Run directory deleted successfully: {run_dir}')
        else:
            self.logger.info(f'Run directory already deleted: {run_dir}')

    def _get_next_run_name(self, experiment_name: str) -> str:
        """
        Generates the next run name for a given experiment.

        Args:
            experiment_name (str): The name of the experiment for which the next run name is being generated.

        Returns:
            str: A unique run name in the format "<experiment_name>-<next_run_number>".
        """
        runs = self.model_serving.search_runs_by_name(
            experiment_names=[experiment_name], order_by=['start_time desc']
        )

        next_run_number = len(runs) + 1
        return f'{experiment_name}-{next_run_number}'

    def _generate_artifacts(self, data: TrainModelResult) -> TrainModelResult:
        """
        Generates and organizes artifacts related to the training process, such as reports and data files.

        Args:
            data: The training model result containing the datasets, model, and parameters.

        Returns:
            The updated result object with paths to the generated artifacts.

        Raises:
            FileNotFoundError: If the reports directory or header.html file does not exist.
            ValueError: If run_name is not set.
        """
        # Validate that run_name is set
        if not data.run_name:
            error_msg = 'run_name must be set before generating artifacts'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        reference_data, current_data = self._init_artifacts_data(data)
        base_path = self._get_reports_directory()

        # Validate that reports directory exists
        if not path.exists(base_path):
            error_msg = f'Reports directory does not exist: {base_path}'
            self.logger.error(error_msg)
            raise FileNotFoundError(error_msg)

        data.run_dir = self._create_run_directory(base_path, data.run_name)
        header_file_path = path.join(base_path, 'header.html')

        # Validate that header.html exists
        if not path.exists(header_file_path):
            error_msg = f'Header file does not exist: {header_file_path}'
            self.logger.error(error_msg)
            raise FileNotFoundError(error_msg)

        self._setup_run_directory(data.run_dir, header_file_path)
        return self._generate_report(reference_data, current_data, data)

    def _save_run(self, data: TrainModelResult):
        """
        Logs the details of a machine learning run, including parameters, metrics, models, and artifacts,
        to the Sientia tracking system.

        Args:
            data: The training model result containing the datasets, model, parameters,
                  and evaluation metrics.

        Raises:
            ValueError: If required metrics or artifacts are missing.
            Exception: If MLflow logging fails for any reason.
        """
        # Validate that required artifacts exist before attempting to log
        if not data.report_path or not path.exists(data.report_path):
            error_msg = f'Report file does not exist: {data.report_path}'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        if not data.train_data_path or not path.exists(data.train_data_path):
            error_msg = f'Training data file does not exist: {data.train_data_path}'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        if not data.test_data_path or not path.exists(data.test_data_path):
            error_msg = f'Test data file does not exist: {data.test_data_path}'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        # Validate that metrics are present
        if data.mse_val is None or data.r2_val is None or data.mae_val is None:
            error_msg = 'One or more metrics (MSE, R2, MAE) are None'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        # Prepare parameters
        train_test_split = f'{data.params.train_size}-{100 - data.params.train_size}'

        interval_strs = [
            (str(interval[0]), str(interval[1]))
            for interval in (data.params.removed_intervals or [])
        ]

        # Set experiment and create run
        self.model_serving.set_experiment(data.params.experiment_name)

        with self.model_serving.save_experiment(
            run_name=data.run_name, description=data.params.experiment_name
        ):
            # Log model parameters
            self.model_serving.log_param('model_type', 'Linear Regression')
            self.model_serving.log_param('target_variable', data.params.target_variable)
            self.model_serving.log_param('input_variables', data.params.variable_columns)
            self.model_serving.log_param('lag_train', data.params.lag_train)
            self.model_serving.log_param('lag_val', data.params.lag_val)
            self.model_serving.log_param('ma', data.params.window)
            self.model_serving.log_param('low_lim', data.params.low_lim)
            self.model_serving.log_param('upp_lim', data.params.upp_lim)
            self.model_serving.log_param('normalized', data.scaler_dict)
            self.model_serving.log_param('ar', data.params.include_ar)
            self.model_serving.log_param('Train_test_split', train_test_split)
            self.model_serving.log_param('Removed_intervals', interval_strs)
            self.model_serving.log_param('Retrain', False)

            # Log evaluation metrics
            self.model_serving.log_metric('MSE', data.mse_val)
            self.model_serving.log_metric('R2', data.r2_val)
            self.model_serving.log_metric('MAE', data.mae_val)

            # Log models
            self.model_serving.log_model(data.process_data, 'data_model')
            self.model_serving.log_model(data.regr, 'prediction_model')

            # Log artifacts
            self.model_serving.log_artifact(data.report_path)
            self.model_serving.log_artifact(data.train_data_path)
            self.model_serving.log_artifact(data.test_data_path)

    def _init_artifacts_data(self, data: TrainModelResult) -> tuple[pd.DataFrame, pd.DataFrame]:
        """
        Prepares the reference and current datasets for artifact generation.

        Args:
            data: The training model result containing the datasets and model.

        Returns:
            tuple: A tuple containing:
                - reference_data: The training dataset with predictions added.
                - current_data: The testing dataset with predictions added.

        Raises:
            ValueError: If training or test datasets are empty or invalid.
            AttributeError: If required attributes are missing from the data object.
        """
        # Validate that required DataFrames are not empty
        # Note: x_train, y_train, x_test, y_test, and regr are required fields in TrainModelResult
        # so we only check if they are empty, not None
        if data.x_train.empty:
            error_msg = 'Training features (x_train) are empty'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        if data.y_train.empty:
            error_msg = 'Training target (y_train) is empty'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        if data.x_test.empty:
            error_msg = 'Test features (x_test) are empty'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        if data.y_test.empty:
            error_msg = 'Test target (y_test) is empty'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        # Validate that predictions exist (y_pred is optional, so check for None)
        if data.y_pred is None:
            error_msg = 'Test predictions (y_pred) are None'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        # Prepare reference data (training set)
        reference_data = pd.concat([data.x_train, data.y_train], axis=1)
        reference_data = reference_data.rename(columns={data.params.target_variable: 'target'})
        reference_data['prediction'] = data.regr.predict(data.x_train)

        # Prepare current data (test set)
        current_data = pd.concat([data.x_test, data.y_test], axis=1)
        current_data = current_data.rename(columns={data.params.target_variable: 'target'})
        current_data['prediction'] = data.y_pred

        return reference_data, current_data

    def _create_run_directory(self, base_path: str, run_name: str) -> str:
        """
        Creates a directory inside the 'reports' folder with the run name and a timestamp.

        Uses microsecond precision in timestamp to minimize collision probability
        in high-concurrency scenarios.

        Args:
            base_path (str): The path to the 'reports' folder.
            run_name (str): The name of the run.

        Returns:
            str: The path to the created directory.

        Raises:
            PermissionError: If there are insufficient permissions to create the directory.
            OSError: If directory creation fails for any other reason.
        """
        # Use microsecond precision to reduce collision probability
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S_%f')
        run_dir = path.join(base_path, f'{run_name}_{timestamp}')

        try:
            makedirs(run_dir, exist_ok=True)
            return run_dir
        except PermissionError as e:
            error_msg = f'Permission denied when creating directory: {run_dir}'
            self.logger.error(error_msg)
            raise PermissionError(error_msg) from e
        except OSError as e:
            error_msg = f'Failed to create directory {run_dir}: {str(e)}'
            self.logger.error(error_msg)
            raise OSError(error_msg) from e

    def _setup_run_directory(self, run_dir: str, header_file_path: str):
        """
        Creates empty files and copies a header file into the specified run directory.

        Note: Lock removed as each run has its own unique directory, so no synchronization
        is needed between different runs. File operations within the same directory are
        atomic at the OS level.

        Args:
            run_dir (str): The path to the run directory where the files will be created.
            header_file_path (str): The path to the header.html file to be copied.

        Raises:
            FileNotFoundError: If the header file does not exist.
            PermissionError: If there are insufficient permissions to create files.
            OSError: If file creation or copying fails for any other reason.
        """
        empty_files = ['data_drift.html', 'data_quality.html', 'regression.html']

        try:
            # Create empty placeholder files
            for file_name in empty_files:
                file_path = path.join(run_dir, file_name)
                with open(file_path, 'w'):
                    pass  # Create empty file

            # Copy header file to run directory
            header_dest = path.join(run_dir, 'header.html')
            shutil.copy(header_file_path, header_dest)
        except FileNotFoundError as e:
            error_msg = f'Header file not found: {header_file_path}'
            self.logger.error(error_msg)
            raise FileNotFoundError(error_msg) from e
        except PermissionError as e:
            error_msg = f'Permission denied when setting up directory: {run_dir}'
            self.logger.error(error_msg)
            raise PermissionError(error_msg) from e
        except OSError as e:
            error_msg = f'Failed to setup run directory {run_dir}: {str(e)}'
            self.logger.error(error_msg)
            raise OSError(error_msg) from e

    def _generate_report(
        self, reference_data: pd.DataFrame, current_data: pd.DataFrame, data: TrainModelResult
    ) -> TrainModelResult:
        """
        Generates a comprehensive report summarizing data quality, data drift, and regression analysis.

        Args:
            reference_data (pd.DataFrame): The training dataset with predictions added.
            current_data (pd.DataFrame): The testing dataset with predictions added.
            data: The training model result containing the datasets, model, and parameters.

        Returns:
            The updated result object with paths to the generated report and data files.

        Raises:
            ValueError: If data conversion to float64 fails or DataFrames are invalid.
            PermissionError: If there are insufficient permissions to write files.
            OSError: If file writing fails for any other reason.
        """
        try:
            # Convert data to float64 for report generation
            # This may raise ValueError if data contains non-numeric values
            reference_data_float = reference_data.astype(np.float64)
            current_data_float = current_data.astype(np.float64)

            # Initialize report generator
            report = Reports(
                reference_data=reference_data_float,
                current_data=current_data_float,
                base_path=data.run_dir,
            )

            # Generate report sections
            report.add_data_quality_section(columns=data.params.variable_columns + ['target'])
            report.add_data_drift_section(columns=data.params.variable_columns + ['target'])
            report.add_regression_section()

            # Validate that run_dir is set (should be set by _create_run_directory)
            if not data.run_dir:
                error_msg = 'run_dir is not set after directory creation'
                self.logger.error(error_msg)
                raise ValueError(error_msg)

            # Save HTML report
            data.report_path = path.join(data.run_dir, 'report.html')
            report.save_all_sections_html(data.report_path)

            # Save training data CSV
            data.train_data_path = path.join(data.run_dir, 'train_data.csv')
            reference_data.to_csv(data.train_data_path, index=False)

            # Save test data CSV
            data.test_data_path = path.join(data.run_dir, 'test_data.csv')
            current_data.to_csv(data.test_data_path, index=False)

            return data
        except ValueError as e:
            error_msg = f'Failed to convert data to float64 for report generation: {str(e)}'
            self.logger.error(error_msg)
            raise ValueError(error_msg) from e
        except PermissionError as e:
            error_msg = f'Permission denied when writing report files to: {data.run_dir}'
            self.logger.error(error_msg)
            raise PermissionError(error_msg) from e
        except OSError as e:
            error_msg = f'Failed to generate report in {data.run_dir}: {str(e)}'
            self.logger.error(error_msg)
            raise OSError(error_msg) from e

    def _get_reports_directory(self) -> str:
        """
        Get the absolute path to the reports directory.

        Returns:
            str: Absolute path to model_manager/reports directory.
        """
        # Get the directory where this file is located (model_manager/utils/repository/)
        current_file_dir = path.dirname(path.abspath(__file__))
        # Navigate up to model_manager/ and then to reports/
        model_manager_dir = path.dirname(path.dirname(current_file_dir))
        reports_dir = path.join(model_manager_dir, 'reports')
        return reports_dir