import asyncio
import logging.config
from typing import Any, cast
import uuid

from temporalio import workflow
from temporalio.worker import Worker
from temporalio.worker.workflow_sandbox import (
    SandboxRestrictions,
    SandboxedWorkflowRunner,
)
import yaml

with workflow.unsafe.imports_passed_through():

    from phederation.storage.tasks import StorageBackgroundTask
    from phederation.actions.workflows import (
        ActivityDeliveryWorkflow,
        DeliveryWorkerTask,
        UpdateNodeInfoTask,
        KeyRotationTask,
    )
    from phederation.federation.server import ActivityPubServer
    from phederation.utils.settings import PhedSettings
    from phederation.actions.workflows import UpdateNodeInfoWorkflow, KeyRotationWorkflow

interrupt_event = asyncio.Event()


async def main_worker(settings_file: str | None = None, logging_file: str | None = None) -> None:
    settings_file = settings_file or "phederation.yaml"
    settings: PhedSettings = PhedSettings(yaml_file=settings_file)
    if not settings.worker.enabled or not settings.worker.start_own_worker:
        print("Worker disabled in settings, shutting down.")
        return

    server: ActivityPubServer = ActivityPubServer(settings=settings)

    if logging_file is None:
        with open(settings_file) as f:
            settings_config = cast(dict[str, Any], yaml.safe_load(f))
            logging_config = cast(dict[str, Any], settings_config["logging"])
    else:
        with open(logging_file) as f:
            logging_config = cast(dict[str, Any], yaml.safe_load(f))
    logging.config.dictConfig(config=logging_config)

    await server.initialize()

    if not server.worker_client:
        print("Worker could not be started: Client not connected.")
        return

    storage_task = StorageBackgroundTask()
    upgrade_node_info_task = UpdateNodeInfoTask(settings.federation.update_nodeinfo_period)
    delivery_worker = DeliveryWorkerTask(settings.federation.delivery_queue_polling_period)
    key_rotation_task = KeyRotationTask(enabled=settings.security.key_rotation_enabled, update_period=1000)

    await storage_task.initialize(server=server)
    await upgrade_node_info_task.initialize(server=server)
    await delivery_worker.initialize(server=server)
    await key_rotation_task.initialize(server=server)

    # Run a worker for the workflow
    async with Worker(
        server.worker_client,
        task_queue=server.settings.worker.tasks_queue_name,
        identity=f"gunicorn-worker-{str(uuid.uuid4())[:8]}",
        workflows=[UpdateNodeInfoWorkflow, ActivityDeliveryWorkflow, KeyRotationWorkflow],
        workflow_runner=SandboxedWorkflowRunner(
            restrictions=SandboxRestrictions.default.with_passthrough_modules("pydantic", "annotated_types", "pydantic_core")
        ),
        activities=[
            upgrade_node_info_task.update_node_info,
            delivery_worker.poll_delivery_queue,
            delivery_worker.prepare_inboxes_for_delivery,
            delivery_worker.deliver_activity,
            key_rotation_task.rotate_keys,
            key_rotation_task.announce_key_rotation,
        ],
    ):
        # Wait until interrupted
        print("Worker started, ctrl+c to exit")
        _ = await interrupt_event.wait()
        print("Shutting down")


if __name__ == "__main__":
    loop = asyncio.new_event_loop()
    try:
        loop.run_until_complete(main_worker())
    except KeyboardInterrupt:
        interrupt_event.set()
        loop.run_until_complete(loop.shutdown_asyncgens())
