import asyncio
from datetime import timedelta
from http import HTTPStatus
from typing import override

from temporalio import workflow
from temporalio.client import (
    Client,
    Schedule,
    ScheduleActionStartWorkflow,
    ScheduleIntervalSpec,
    ScheduleSpec,
    ScheduleState,
)
from temporalio.common import RetryPolicy

from phederation.actions.tasks import KeyRotationTask, UpdateNodeInfoTask


with workflow.unsafe.imports_passed_through():

    from .tasks import (
        CanBeScheduled,
        DeliveryWorkerTask,
        DeliveryWorkflowParams,
    )
    from phederation.utils.settings import PhedSettings
    from phederation.utils.logging import configure_logger
    from phederation.utils.base import ObjectId
    from phederation.federation.delivery import ActivityDelivery, DeliveryResult
    from phederation.utils.base import ObjectId
    from phederation.utils.settings import FederationSettings


@workflow.defn
class UpdateNodeInfoWorkflow(CanBeScheduled):
    """Workflow to periodically update the NodeInfo of a running instance.

    It is usually used in a temporalio scheduler, and started by the `UpdateNodeInfoTask`.
    """

    @staticmethod
    @override
    async def schedule(settings: PhedSettings, client: Client):
        return await client.create_schedule(
            "workflow-schedule-UpdateNodeInfoTask",
            Schedule(
                action=ScheduleActionStartWorkflow(
                    UpdateNodeInfoWorkflow.run,
                    id="schedules-workflow-UpdateNodeInfoTask",
                    task_queue=settings.worker.tasks_queue_name,
                ),
                spec=ScheduleSpec(intervals=[ScheduleIntervalSpec(every=timedelta(seconds=settings.federation.update_nodeinfo_period))]),
                state=ScheduleState(note="Periodic update of the instance node info data."),
            ),
        )

    @workflow.run
    async def run(
        self,
    ) -> ObjectId:
        """Run the workflow.

        Returns:
            ObjectId: The id of the NodeInfo document.
        """
        retry_policy = RetryPolicy(initial_interval=timedelta(seconds=1), maximum_attempts=3)
        document_id = await workflow.execute_activity_method(
            UpdateNodeInfoTask.update_node_info,
            retry_policy=retry_policy,
            schedule_to_close_timeout=timedelta(seconds=20),
        )
        print(f"UpdateNodeInfoWorkflow: Updated instance node information.")
        return document_id


@workflow.defn
class KeyRotationWorkflow(CanBeScheduled):
    """Workflow to periodically rotate keys of a running instance.

    It is usually used in a temporalio scheduler, and started by the `KeyRotationTask`.
    """

    @staticmethod
    @override
    async def schedule(settings: PhedSettings, client: Client):
        return await client.create_schedule(
            "workflow-schedule-KeyRotationTask",
            Schedule(
                action=ScheduleActionStartWorkflow(
                    KeyRotationWorkflow.run,
                    settings.federation,
                    id="schedules-workflow-KeyRotationTask",
                    task_queue=settings.worker.tasks_queue_name,
                ),
                spec=ScheduleSpec(intervals=[ScheduleIntervalSpec(every=timedelta(seconds=settings.security.key_rotation_period))]),
                state=ScheduleState(note="Periodic update of the instance key data."),
            ),
        )

    @workflow.run
    async def run(self, federation_settings: FederationSettings):
        """Run the workflow."""
        retry_policy_rotate = RetryPolicy(initial_interval=timedelta(seconds=10), maximum_attempts=1)
        retry_policy_announce = RetryPolicy(
            initial_interval=timedelta(seconds=5), backoff_coefficient=2.0, maximum_attempts=federation_settings.max_retries
        )

        logger = configure_logger(name=__name__, prefix=federation_settings.logging_prefix)

        logger.debug("Rotating keys")
        key_rotation_parameters = await workflow.execute_activity_method(
            KeyRotationTask.rotate_keys,
            retry_policy=retry_policy_rotate,
            schedule_to_close_timeout=timedelta(seconds=30),
        )
        logger.debug(f"Announcing rotated keys: {key_rotation_parameters}")
        for params in key_rotation_parameters:
            print(params)
            _ = await workflow.execute_activity_method(
                KeyRotationTask.announce_key_rotation,
                params,
                retry_policy=retry_policy_announce,
                schedule_to_close_timeout=timedelta(seconds=30),
            )
        logger.info(f"KeyRotationWorkflow: Rotated instance keys.")


@workflow.defn
class ActivityDeliveryWorkflow(CanBeScheduled):

    @staticmethod
    @override
    async def schedule(settings: PhedSettings, client: Client):
        handle = await client.create_schedule(
            "workflow-schedule-DeliveryWorkerTask",
            Schedule(
                action=ScheduleActionStartWorkflow(
                    ActivityDeliveryWorkflow.run,
                    settings.federation,
                    id="schedules-workflow-DeliveryWorkerTask",
                    task_queue=settings.worker.tasks_queue_name,
                ),
                spec=ScheduleSpec(intervals=[ScheduleIntervalSpec(every=timedelta(seconds=settings.federation.delivery_queue_polling_period))]),
                state=ScheduleState(note="Periodic polling of delivery queue for new messages to send."),
            ),
        )
        return handle

    @workflow.run
    async def run(self, settings_federation: FederationSettings):
        logger = configure_logger(name=__name__, prefix=settings_federation.logging_prefix)
        logger.debug(f"ActivityDeliveryWorkflow: Polling delivery queue...")

        retry_policy = RetryPolicy(initial_interval=timedelta(seconds=3), maximum_attempts=1)
        params = await workflow.execute_activity_method(
            DeliveryWorkerTask.poll_delivery_queue,
            settings_federation,
            retry_policy=retry_policy,
            schedule_to_close_timeout=timedelta(seconds=settings_federation.delivery_timeout),
        )
        if not params:
            logger.debug(f"ActivityDeliveryWorkflow: Nothing to do.")
            return

        logger.info(f"ActivityDeliveryWorkflow: Starting delivery to {len(params.recipients)} recipients...")
        # first, make sure that all bcc and bto recipients are actually removed
        params.activity = await ActivityDelivery.remove_hidden_recipients(params.activity)

        # then prepare all inboxes for all recipients, deduplicating the list in the process
        retry_policy = RetryPolicy(initial_interval=timedelta(seconds=3), maximum_attempts=params.settings_federation.max_retries)
        recipient_inboxes = await workflow.execute_activity_method(
            DeliveryWorkerTask.prepare_inboxes_for_delivery,
            params,
            retry_policy=retry_policy,
            schedule_to_close_timeout=timedelta(seconds=params.settings_federation.delivery_timeout),
        )
        logger.debug(f"Prepared inboxes of {len(recipient_inboxes)} recipients")

        # finally, send off the activity to each recipient
        retry_policy = RetryPolicy(
            initial_interval=timedelta(seconds=1), backoff_coefficient=2.0, maximum_attempts=params.settings_federation.max_retries
        )
        handles: list[workflow.ActivityHandle[DeliveryResult]] = []
        for inbox in recipient_inboxes:
            logger.debug(f"Sending activity to {inbox}...")
            handle = workflow.start_activity_method(
                DeliveryWorkerTask.deliver_activity,
                DeliveryWorkflowParams(activity=params.activity, inbox=inbox),
                retry_policy=retry_policy,
                schedule_to_close_timeout=timedelta(seconds=params.settings_federation.delivery_timeout),
                summary="Deliver a given activity to a given inbox",
            )
            handles.append(handle)
        _ = await asyncio.gather(*handles)

        # wait on all send activities to be done
        await workflow.wait_condition(lambda: workflow.all_handlers_finished())

        logger.debug(f"Every message sent, collecting results")
        # collect all results
        delivery_results: list[DeliveryResult] = []
        for h in handles:
            if h.done() and not h.cancelled():
                delivery_results.append(h.result())
            else:
                # context = h.get_context()
                if params.activity.id:
                    delivery_results.append(DeliveryResult(failed=[params.activity.id], status_code=HTTPStatus.INTERNAL_SERVER_ERROR))

        logger.info(f"ActivityDeliveryWorkflow: Delivery done.")
        return delivery_results
