from temporalio import activity, workflow

with workflow.unsafe.imports_passed_through():
    import json
    import traceback
    from typing import Any

    from pandas import DataFrame
    from sientia_do.notifications.handlers import NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from sientia_do.observability.logger import Logger
    from sientia_do.observability.metrics_controller import MetricsController
    from sientia_do.observability.sientia_monitoring import SientiaMonitoring
    from sientia_do.repository.pi_web_api_client import PIWebAPIClient

    from laborious import metrics


PI_WEB_API_PREDICTION_ERROR_CONFIDENCE = 13


class API(SientiaMonitoring):
    """
    PI Web API operations for writing prediction data to PI Web API.

    This class provides Temporal activities for interacting with the PI Web API
    to write prediction and confidence values to industrial systems. It handles
    error scenarios gracefully by setting error confidence values and sending
    notifications when write operations fail.

    The class implements comprehensive error handling for both prediction and
    confidence value writes, ensuring that partial failures are properly
    reported and handled.
    """

    def __init__(
        self,
        base_url: str,
        auth_type: str,
        auth_token: str,
        logger: Logger,
        notification_handler: NotificationHandler,
        metrics_controller: MetricsController,
    ) -> None:
        """
        Initialize API activity with PI Web API client.

        Args:
            base_url (str): Base URL of the PI Web API server
            auth_type (str): Authentication type ('basic' or 'bearer')
            auth_token (str): Authentication token
            logger (Logger): Logger instance for operation logging
            notification_handler (NotificationHandler): Handler for system notifications
            metrics_controller (MetricsController): Controller for metrics collection
        """
        SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
        self.pi_web_api_client = PIWebAPIClient(
            base_url=base_url,
            auth_config={
                'type': auth_type,
                'token': auth_token,
            },
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

    def close(self) -> None:
        """
        Close the PI Web API client and shutdown monitoring services.

        This method properly closes all connections and resources associated
        with the PI Web API client and monitoring services.
        """
        self.pi_web_api_client.close()
        SientiaMonitoring.shutdown(self)

    async def process_pi_web_api_response(
        self,
        response_data: dict[str, Any],
        tags: dict[str, str],
        core_labels: dict[str, str],
        metadata: dict[str, Any],
    ) -> tuple[int, str]:
        """
        Process the response data from PI Web API write operation.

        Validates that all tags were successfully written, emits metrics for each tag
        (success or error), and returns the appropriate prediction confidence value.
        Sets error confidence if any tag write fails or if the number of written tags
        doesn't match the expected count.

        Args:
            - response_data (dict[str, Any]): The response data from the PI Web API write operation.
            - tags (dict[str, str]): The tags that were written to the PI Web API.
            - core_labels (dict[str, str]): The core labels of the workflow execution.
            - metadata (dict[str, Any]): The metadata of the workflow execution.
        Returns:
            int: Prediction confidence value (0 for success, 13 for errors)
        """

        # Convert tags from name:webid to webid:name
        tags = {w: t for t, w in tags.items()}

        tag_names = list[str](tags.values())

        confidence = 0

        message = ''

        # Evaluate response for each tag
        written_tags = []
        response_items = response_data.get('Items', [])
        for item in response_items:
            web_id = item.get('WebId')
            if not web_id:
                self.error('The response did not contain some WebIds', metadata)
                continue
            errors = item.get('Errors', [])
            tag_name = tags.get(web_id)
            if not tag_name:
                self.error(
                    f'The response did not contain the tag name for WebId {web_id}', metadata
                )
                continue
            if errors:
                self.error(
                    f'Error writing tag {tag_name}:{web_id} to PI Web API: {errors}', metadata
                )
                await self.emit_metric(
                    metric_object=metrics.PI_WEB_API_PREDICTION_WRITTEN_ERROR_COUNT,
                    tags={
                        **core_labels,
                        'tag_name': tag_name,
                    },
                )
                confidence = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
            else:
                await self.emit_metric(
                    metric_object=metrics.PI_WEB_API_PREDICTION_WRITTEN_COUNT,
                    tags={
                        **core_labels,
                        'tag_name': tag_name,
                    },
                )
                written_tags.append(tag_name)

        if len(written_tags) != len(tag_names):
            message = f'The number of written tags does not match the number of tag names: Expected {tag_names} tags, but {written_tags} tags were written.'

            self.error(
                f'{message}\nResponse:\n {json.dumps(response_data, indent=4)}\nTags:\n {json.dumps(tags, indent=4)}',
                metadata,
            )

            await self.send_notification_async(
                metadata=metadata,
                notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
                message=f'The number of written tags does not match the number of tag names: Expected {tag_names} tags, but {written_tags} tags were written.\nResponse:\n {json.dumps(response_data, indent=4)}\nTags:\n {json.dumps(tags, indent=4)}',
                block='write_pi_web_api_data',
                level=NotificationLevel.ERROR,
            )
            confidence = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE

        return confidence, message

    @activity.defn(name='write_pi_web_api_data')
    async def write_pi_web_api_data(self, input_data: dict[str, Any]) -> dict[Any, Any]:
        """
        Write prediction and confidence data to PI Web API.

        Writes prediction values and confidence scores to PI Web API using configured
        web IDs. Processes responses to validate writes and emit metrics. Handles errors
        gracefully by setting error confidence values when writes fail and sending
        notifications for both prediction and confidence write errors.

        Args:
            input_data (dict[str, Any]): The input data containing:
                - metadata (dict[str, Any]): Workflow execution metadata
                - pi_web_api_output_config (dict[str, Any]): PI Web API configuration with:
                    - endpoint (str): PI Web API endpoint URL
                    - prediction_tags (dict[str, str]): Mapping of tag names to web IDs for predictions
                    - confidence_tags (dict[str, str]): Mapping of tag names to web IDs for confidence
                - data (dict[str, Any]): Prediction data, its a dataframe converted to dict.
        Returns:
            dict[Any, Any]: Data dictionary with potentially modified confidence values
                If prediction write fails, prediction_confidence is set to error value (13)
        """
        metadata = input_data['metadata']
        data = DataFrame(input_data['data'])
        pi_web_api_output_config = input_data['pi_web_api_output_config']

        self.info(f'Writing data to PI Web API... config: {pi_web_api_output_config}', metadata)

        raw_prediction_tags = pi_web_api_output_config['prediction_tags']
        raw_confidence_tags = pi_web_api_output_config['confidence_tags']
        prediction_tags = list[str](raw_prediction_tags.values())
        confidence_tags = list[str](raw_confidence_tags.values())

        core_labels = self.get_core_labels(metadata)

        prediction_value = data.head(1)['prediction'].values[0]
        confidence_value = data.head(1)['prediction_confidence'].values[0]

        try:
            prediction_response = await self.pi_web_api_client.write_value(
                web_ids=prediction_tags,
                value={
                    'Timestamp': data.head(1)['timestamp'].values[0],
                    'Value': prediction_value,
                },
                metadata=metadata,
            )

            confidence, message = await self.process_pi_web_api_response(
                response_data=prediction_response,
                tags=raw_prediction_tags,
                core_labels=core_labels,
                metadata=metadata,
            )

            data['prediction_confidence'] = confidence
            data['comments'] = message

        except Exception as e:
            trace = traceback.format_exc()
            await self.send_notification_async(
                metadata=metadata,
                notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
                message=f'Error writing prediction data to PI Web API: {e}\n Tags: {raw_prediction_tags}',
                block='write_pi_web_api_data',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )

            data['prediction_confidence'] = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
            data['comments'] = str(e)

            return data.to_dict()

        try:
            confidence_response = await self.pi_web_api_client.write_value(
                web_ids=confidence_tags,
                value={
                    'Timestamp': data.head(1)['timestamp'].values[0],
                    'Value': confidence_value,
                },
                metadata=metadata,
            )

            await self.process_pi_web_api_response(
                response_data=confidence_response,
                tags=raw_confidence_tags,
                core_labels=core_labels,
                metadata=metadata,
            )

        except Exception as e:
            trace = traceback.format_exc()
            await self.send_notification_async(
                metadata=metadata,
                notification_id='WRITE_PI_WEB_API_CONFIDENCE_ERROR',
                message=f'Error writing confidence data to PI Web API: {e}\n Tags: {raw_confidence_tags}',
                block='write_pi_web_api_data',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )

        return data.to_dict()