from temporalio import workflow, client
from temporalio.worker import Worker
import sys

with workflow.unsafe.imports_passed_through():
    import os
    import asyncio
    from laborious.workflows.predictions_batch import PredictionsBatch
    from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
    from laborious.workflows.sub_workflows.format_and_export_prediction import \
        FormatAndExportPrediction
    from laborious.activities.activities import Activities
    from laborious.utils.logger import get_logger
    from laborious.utils.connectors_config import (
        build_postgres_config,
        build_mlflow_config,
        build_opc_config
    )
    from sientia_do.notifications.handlers import NotificationHandler


async def main():
    host = os.getenv('TEMPORAL_HOST', 'localhost:7233')
    logger = get_logger(__name__)

    logger.info('Starting Worker...')

    logger.info('Starting Notification Handler...')

    notification_handler = NotificationHandler(
        servers=os.getenv('KAFKA_SERVERS', 'http://localhost:9092'),
        logger=logger,
        project_name=os.getenv('PROJECT_NAME', 'laborious'),
        pipeline_name='-',
        trigger_name='-',
        model_name='-',
        model='-'
    )

    logger.info('Starting Activities...')

    activities = Activities(
        postgres_config=build_postgres_config(),
        mlflow_config=build_mlflow_config(),
        opc_config=build_opc_config(),
        logger=logger,
        notification_handler=notification_handler
    )

    logger.info('Starting Temporal Client...')

    temporal_client = await client.Client.connect(
        target_host=host,
        namespace=os.getenv('TEMPORAL_NAMESPACE', 'default')
    )

    logger.info('Starting Workers...')

    workers = [
        Worker(
            temporal_client,
            task_queue='predictions-queue',
            workflows=[PredictionsBatch, PredictionProcess,
                       FormatAndExportPrediction],
            activities=[
                # Base
                activities.prepare_activity,
                # MLFlow
                activities.request_predict,
                activities.request_transform,
                # Gates
                activities.input_gate,
                activities.mlflow_response_gate,
                activities.mlflow_content_gate,
                activities.format_prediction,
                activities.format_default_prediction,
                activities.get_last_timestamp,
                # OPC
                activities.write_opc_data,
                # Postgres
                activities.load_custom_query,
                activities.repeat_last_prediction,
                activities.export_data_to_postgres
            ]
        )
    ]

    handlers = []
    for w in workers:
        handlers.append(w.run())

    logger.info('Workers started successfully')

    try:
        # This will run the workers and wait for them to complete.
        # If an exception occurs in any of the worker handlers, it will be propagated here.
        await asyncio.gather(*handlers)
    except BaseException as e:
        logger.error("An unhandled exception occurred: %s", e, exc_info=True)
    finally:
        if notification_handler:
            notification_handler.shutdown()
        if activities:
            activities.shutdown()
        # Exit with a non-zero status code to indicate failure to Kubernetes
        sys.exit(1)

if __name__ == '__main__':
    asyncio.run(main())