from temporalio import workflow

with workflow.unsafe.imports_passed_through():
    from laborious.activities.activities import Activities
    from typing import Any
    from sientia_do.temporal.policies import retry_policy
    from datetime import timedelta


@workflow.defn(name="prediction_process")
class PredictionProcess():
    """
    Core prediction processing workflow for the Laborious system.

    This workflow implements the complete ML model inference pipeline, handling
    data quality validation, MLFlow model interactions, and prediction processing.
    It serves as the central orchestrator for all prediction operations and ensures
    data quality throughout the entire process.

    The workflow implements a robust data processing pipeline with:
    - Data quality validation using configurable filters
    - MLFlow model transformation and prediction
    - Response validation and quality assurance
    - Flexible decision path handling
    - Comprehensive error handling and retry policies

    Workflow Execution:
        1. Timestamp Retrieval: Gets last processed timestamp for incremental processing
        2. Input Data Gate: Applies data quality filters
        3. Path Decision: Determines processing path based on filter results
        4. MLFlow Transform: Requests data transformation using MLFlow models
        5. Response Validation: Filters transform responses for quality assurance
        6. MLFlow Prediction: Executes prediction using transformed data
        7. Content Validation: Filters prediction responses for final quality check
        8. Export Delegation: Delegates to FormatAndExportPrediction workflow
    """

    @workflow.run
    async def run(self, input_data: dict[str, Any]):
        """
        Execute the prediction process workflow.

        This method orchestrates the complete prediction processing pipeline by:
        1. Retrieving the last processed timestamp for incremental processing
        2. Applying data quality filters to validate input data
        3. Executing MLFlow model transformation and prediction
        4. Validating all responses for quality assurance
        5. Delegating to export workflow for data persistence

        The method implements comprehensive error handling and ensures all
        data quality requirements are met before proceeding with ML operations.

        Args:
            input_data: Complete configuration for the prediction process
                Required keys:
                    - metadata (dict): Workflow execution metadata
                    - data (dict): Input data for prediction processing
                    - schema (dict): Data schema definition
                    - table_name (str): Target table for predictions
                    - model_id (str): ML model identifier
                    - model_name (str): ML model name
                    - input_filters (dict): Data quality filters
                    - mlflow_transform_filters (dict): MLFlow transform filters
                    - mlflow_predict_filters (dict): MLFlow prediction filters
                    - model_retention (int): Model retention period in minutes
                    - path_priority (list[str]): Decision path priority configuration
                    - opc_output_config (dict): OPC server export configuration

        Returns:
            None: The workflow completes successfully when export workflow finishes

        Raises:
            Exception: If any required parameters are missing or if the workflow fails
                during data processing, MLFlow operations, or workflow delegation

        """

        metadata = input_data['metadata']
        data = input_data['data']
        model_id = input_data['model_id']
        model_name = input_data['model_name']
        model_config = input_data.get('model_config', {})

        # Get last timestamp for incremental processing
        last_timestamp = await workflow.execute_local_activity_method(
            Activities.get_last_timestamp,
            {
                **metadata,
                'data': data
            },
            retry_policy=retry_policy,
            start_to_close_timeout=timedelta(minutes=1),
        )

        # Apply input data quality gates
        gate_input = {
            **metadata,
            'filters': input_data['input_filters'],
            'data': data,
            'path_priority': input_data['path_priority']
        }

        path_flag, confidence, comment = await workflow.execute_local_activity_method(
            Activities.input_gate,
            gate_input,
            retry_policy=retry_policy,
            start_to_close_timeout=timedelta(minutes=1),
        )

        # Handle path decision based on filter results
        if await self.path_flag_handler(
            data, path_flag, input_data, confidence, last_timestamp, comment
        ):
            return

        # Request MLFlow model transformation
        response_data = await workflow.execute_local_activity_method(
            Activities.request_transform,
            {
                **metadata,
                'data': data,
                'model_name': model_name,
                'model_config': model_config
            },
            retry_policy=retry_policy,
            start_to_close_timeout=timedelta(minutes=5),
        )

        # Validate MLFlow transform response
        path_flag, confidence, comment = await workflow.execute_local_activity_method(
            Activities.mlflow_response_gate,
            {
                **metadata,
                'filters': input_data['mlflow_transform_filters'],
                'data': response_data,
                'type': 'transform',
                'path_priority': input_data['path_priority']
            },
            retry_policy=retry_policy,
            start_to_close_timeout=timedelta(minutes=1),
        )

        # Handle path decision based on transform validation
        if await self.path_flag_handler(
            data, path_flag, input_data, confidence, last_timestamp, comment
        ):
            return

        transformed_data = response_data['content']

        path_flag, confidence, comment = await workflow.execute_local_activity_method(
            Activities.mlflow_content_gate,
            {
                **metadata,
                'filters': input_data['mlflow_transform_filters'],
                'data': transformed_data,
                'type': 'transform',
                'path_priority': input_data['path_priority']
            },
            retry_policy=retry_policy,
            start_to_close_timeout=timedelta(minutes=1),
        )

        if await self.path_flag_handler(
            data, path_flag, input_data, confidence, last_timestamp, comment
        ):
            return

        response_data = await workflow.execute_local_activity_method(
            Activities.request_predict,
            {
                **metadata,
                'data': transformed_data,
                'model_name': model_name,
                'model_config': model_config
            },
            retry_policy=retry_policy,
            start_to_close_timeout=timedelta(minutes=5),
        )

        # Validate MLFlow prediction response
        path_flag, confidence, comment = await workflow.execute_local_activity_method(
            Activities.mlflow_response_gate,
            {
                **metadata,
                'filters': input_data['mlflow_predict_filters'],
                'data': response_data,
                'type': 'predict',
                'path_priority': input_data['path_priority']
            },
            retry_policy=retry_policy,
            start_to_close_timeout=timedelta(minutes=1),
        )

        # Handle path decision based on prediction validation
        if await self.path_flag_handler(
            data, path_flag, input_data, confidence, last_timestamp, comment
        ):
            return

        # Delegate to export workflow for data persistence
        await workflow.execute_child_workflow(
            'format_and_export_prediction',
            {
                'metadata': metadata,
                'path_flag': path_flag,
                'data': response_data['content'],
                'prediction_confidence': confidence,
                'timestamp': last_timestamp,
                'model_id': model_id,
                'model_name': model_name,
                'model_config': model_config,
                'opc_output_config': input_data['opc_output_config'],
                'schema': input_data['schema'],
                'table_name': input_data['table_name'],
                'comment': comment,
                'prediction_store_policy': input_data['prediction_store_policy']
            }
        )

    async def path_flag_handler(self, data: dict, path_flag: str, input_data: dict,
                                confidence: int, last_timestamp: str, comment: str) -> bool:
        """
        Handle path decisions based on filter results and confidence levels.

        This method determines the appropriate action based on the path flag
        returned by data quality filters. It can stop processing, continue,
        or repeat operations based on the configured path priority.

        Args:
            data: Input data for processing
            path_flag: Path decision from filter (STOP, CONTINUE, REPEAT)
            input_data: Complete workflow input configuration
            confidence: Confidence level from filter validation
            last_timestamp: Last processed timestamp
            comment: Additional information about the filter result

        Returns:
            bool: True if processing should stop, False to continue

        Path Handling:
            - STOP: Terminates workflow execution
            - CONTINUE: Proceeds with normal processing
            - REPEAT: Repeats last prediction if available
        """
        metadata = input_data['metadata']

        schema = input_data['schema']
        table_name = input_data['table_name']
        model_id = input_data['model_id']
        model_name = input_data['model_name']
        model_config = input_data.get('model_config', {})

        path_flag = path_flag.upper() if path_flag else ''

        if path_flag == 'STOP':
            # Stop processing and exit workflow
            return True
        elif path_flag == 'REPEAT':
            # Repeat last prediction if available
            await workflow.execute_activity_method(
                Activities.repeat_last_prediction,
                {
                    **metadata,
                    'schema': schema,
                    'table_name': table_name,
                    'model': model_id,
                    'last_timestamp': last_timestamp
                },
                retry_policy=retry_policy,
                start_to_close_timeout=timedelta(minutes=1),
            )
            return True
        elif path_flag == 'CONTINUE':
            # call write workflow
            await workflow.execute_child_workflow(
                'format_and_export_prediction',
                {
                    'metadata': metadata,
                    'path_flag': path_flag,
                    'data': data,
                    'prediction_confidence': confidence,
                    'timestamp': last_timestamp,
                    'model_id': model_id,
                    'model_name': model_name,
                    'model_config': model_config,
                    'schema': schema,
                    'table_name': table_name,
                    'comment': comment,
                    'opc_output_config': input_data['opc_output_config'],
                    'prediction_store_policy': input_data['prediction_store_policy']
                }
            )
            return True

        return False