from temporalio import activity, workflow


with workflow.unsafe.imports_passed_through():
    from logging import Logger
    from sientia_do.notifications.handlers import NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from laborious.activities.base import BaseActivity
    from laborious.utils.repository.opc_repository import OpcRepository
    from typing import Any
    import traceback
    from pandas import DataFrame


class OPC(BaseActivity):
    def __init__(self, opc_servers: dict[str, dict[str, Any]],
                 logger: Logger, notification_handler: NotificationHandler):

        self.logger = logger
        self.notification_handler = notification_handler
        self.opc_servers = opc_servers

        self.opc_repository = {}
        for name, server in opc_servers.items():
            self.opc_repository[name] = OpcRepository(
                name=name,
                url=server['url'],
                logger=self.logger,
                server_uri=server['server_uri'],
                cert_path=server['cert_path'],
                private_key_path=server['private_key_path'],
                server_cert_path=server['server_cert_path'],
                notification_handler=self.notification_handler,
                reconnection_interval=server['reconnection_interval'],
            )
            self.opc_repository[name].connect()

        BaseActivity.__init__(self, logger, notification_handler)

    def write_data(self, server: str, tag: str, data: Any,
                   data_type: str, tag_type: str):
        try:
            self.opc_repository[server].write_data(
                tag, data, data_type)
            self.logger.debug(f"Wrote {tag_type} to {tag}")
        except Exception as e:
            trace = traceback.format_exc()
            self.notification_handler.build_and_send_notification(
                notification_id=f"WRITE_OPC_{tag_type.upper()}_ERROR",
                message=f"Error writing data to OPC server: {e}",
                block="write_opc_data",
                level=NotificationLevel.ERROR,
                attachment_content=trace
            )
            self.logger.error(trace)

    @activity.defn(name='write_opc_data')
    async def write_opc_data(self, input_data: dict[str, Any]):
        """
        Write prediction and confidence data to OPC servers. The two writing
        operations are optional and independent of each other.

        Args:
            input_data (dict[str, Any]): The input data. Contains the following keys:
                - data (dict[str, Any]): The dataframe that contains the data to write
                to the OPC servers.
                - opc_output_config (dict[str, Any]): The OPC writing configuration.
                The keys are the OPC server names and the values contain:
                    prediction_tags (dict[str, Any]): The tags to write to the OPC servers.
                    confidence_tags (dict[str, Any]): The tags to write to the OPC servers.

        Returns:
        """
        self.logger.debug("Writing data to OPC servers...")
        data = DataFrame(input_data['data'])
        opc_output_config = input_data['opc_output_config']
        self.logger.debug(data)

        for server, config in opc_output_config.items():
            if self.opc_repository.get(server) is None:
                self.logger.error(f"OPC server {server} not found")
                continue

            if 'prediction_tags' in config:
                for tag, tag_config in config['prediction_tags'].items():
                    self.write_data(
                        server=server,
                        tag=tag,
                        data=data.head(1)['prediction'].values[0],
                        data_type=tag_config['data_type'],
                        tag_type='prediction'
                    )
            if 'confidence_tags' in config:
                for tag, tag_config in config['confidence_tags'].items():
                    self.write_data(
                        server=server,
                        tag=tag,
                        data=data.head(1)['prediction_confidence'].values[0],
                        data_type=tag_config['data_type'],
                        tag_type='confidence'
                    )

    def shutdown(self):
        for opc in self.opc_repository.values():
            opc.disconnect()