from temporalio import workflow, client
from temporalio.worker import Worker, PollerBehaviorAutoscaling
from temporalio.runtime import Runtime, TelemetryConfig, PrometheusConfig

with workflow.unsafe.imports_passed_through():
    import os
    import sys
    import asyncio
    from laborious.workflows.minimal_retrain import MinimalRetrain
    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.connectors_config import (
        build_postgres_config,
        build_mlflow_config,
        build_opc_config,
        build_mongodb_config
    )
    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    from sientia_do.observability.logger import get_logger
    from laborious import metrics
    from prometheus_client import start_http_server

POD_ID = os.getenv('POD_ID')
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', "9091"))


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

    metadata = {
        'pod_id': POD_ID,
        'model_name': '-',
        'model_id': '-',
        'workflow_name': '-',
        'schedule_name': '-',
    }

    logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata)

    logger.custom_info("Starting prometheus client...", metadata)
    start_prometheus_server()

    logger.custom_info('Starting Notification Handler...', metadata)

    mongo_config = build_mongodb_config()
    notification_handler = NotificationHandler(
        connection_string=mongo_config['connection_string'],
        database=mongo_config['database_name'],
        logger=logger,
        project_name=os.getenv('PROJECT_NAME', 'laborious')
    )

    logger.custom_info('Starting Activities...', metadata)

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

    logger.custom_info('Initializing OPC...', metadata)
    await activities.init_opc()

    logger.custom_info(
        f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)

    new_runtime = Runtime(
        telemetry=TelemetryConfig(
            metrics=PrometheusConfig(
                bind_address=f"0.0.0.0:{SDK_METRICS_PORT}")
        )
    )

    logger.custom_info(f'Starting Temporal Client at {host}...', metadata)

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

    logger.custom_info('Starting Workers...', metadata)

    workers = [
        Worker(
            temporal_client,
            task_queue='minimal_retrain-queue',
            workflows=[MinimalRetrain],
            activities=[
                activities.load_custom_query,
                activities.retrain_model,
                activities.update_production_model,
                activities.export_data_to_postgres
            ],
            max_concurrent_workflow_tasks=50,
            max_concurrent_activities=50,
            max_concurrent_local_activities=50,
            max_cached_workflows=200,
            workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
            activity_task_poller_behavior=PollerBehaviorAutoscaling()
        ),
        Worker(
            temporal_client,
            task_queue='predictions_batch-queue',
            workflows=[PredictionsBatch, PredictionProcess,
                       FormatAndExportPrediction],
            activities=[
                # 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,
                activities.write_metrics
            ],
            max_concurrent_workflow_tasks=50,
            max_concurrent_activities=50,
            max_concurrent_local_activities=50,
            max_cached_workflows=200,
            workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
            activity_task_poller_behavior=PollerBehaviorAutoscaling()
        )
    ]

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

    logger.custom_info('Workers started successfully', metadata)

    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:  # NOSONAR
        logger.custom_error(f"An unhandled exception occurred: {e}", metadata)
    finally:
        if notification_handler:
            notification_handler.shutdown()
        if activities:
            await activities.shutdown()
        # Exit with a non-zero status code to indicate failure to Kubernetes
        metrics.APP_UP.labels(pod_id=POD_ID).set(0)  # Mark app as DOWN
        sys.exit(1)


def start_prometheus_server():
    try:
        port = int(os.getenv("HTTP_METRICS_PORT", 9090))
        start_http_server(port)
        print(f"Prometheus server started on port {port}.")
        metrics.APP_UP.labels(pod_id=POD_ID).set(1)  # Mark app as UP
    except Exception as e:
        print(f"Failed to start Prometheus server: {e}")
        os._exit(1)


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