import asyncio
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from http import HTTPStatus
from typing import Any, cast, override

from temporalio import activity, workflow
from temporalio.client import (
    Schedule,
    ScheduleActionStartWorkflow,
    ScheduleHandle,
    ScheduleIntervalSpec,
    ScheduleSpec,
    ScheduleState,
)
from temporalio.common import RetryPolicy
from temporalio.exceptions import ApplicationError

from phederation.utils.exceptions import DeliveryError

with workflow.unsafe.imports_passed_through():

    from phederation.utils.base import AccessType
    from phederation.utils.base import UrlType
    from phederation.utils.serialization import to_json_string
    from phederation.utils.logging import configure_logger
    from phederation.models.activities import APAnnounce
    from phederation.utils.base import ObjectId
    from phederation.utils.exceptions import ActivityPubException
    from phederation.federation.tasks import PeriodicBackgroundTask
    from phederation.models import dereference
    from phederation.federation.delivery import ActivityDelivery, DeliveryResult
    from phederation.models import ValidCollection
    from phederation.models.activities import APActivity
    from phederation.models.collections import APOrderedCollection
    from phederation.utils.base import ObjectId, collection_id_from_name
    from phederation.utils.exceptions import (
        DiscoveryError,
        ResolverError,
        UserError,
    )
    from phederation.utils.settings import FederationSettings


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

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

    @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:
    """Workflow to periodically rotate keys of a running instance.

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

    @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.")


@dataclass
class PrepareWorkflowParams:
    settings_federation: FederationSettings
    activity: APActivity
    recipients: list[ObjectId]


@dataclass
class DeliveryWorkflowParams:
    activity: APActivity
    inbox: ObjectId


class DeliveryWorkerTask(PeriodicBackgroundTask):
    def __init__(self, update_period: int):
        super().__init__(update_period)
        self.schedule_handle: ScheduleHandle | None = None

    @override
    async def on_startup(self) -> None:
        if not self.server:
            self.logger.warning("Trying to start DeliveryWorkerTask without initialized server.")
            return

        if not self.server.settings.worker.enabled or not self.server.worker_client:
            self.logger.warning(
                "Trying to start DeliveryWorkerTask without enabling temporalio worker. Schedule will not be started, using internal periodic task."
            )
            return

        self.schedule_handle = await self.server.worker_client.create_schedule(
            "workflow-schedule-DeliveryWorkerTask",
            Schedule(
                action=ScheduleActionStartWorkflow(
                    ActivityDeliveryWorkflow.run,
                    self.server.settings.federation,
                    id="schedules-workflow-DeliveryWorkerTask",
                    task_queue=self.server.settings.worker.tasks_queue_name,
                ),
                spec=ScheduleSpec(
                    intervals=[ScheduleIntervalSpec(every=timedelta(seconds=self.server.settings.federation.delivery_queue_polling_period))]
                ),
                state=ScheduleState(note="Periodic polling of delivery queue for new messages to send."),
            ),
        )

        self.logger.info(f"Scheduled updating node info data every {self.server.settings.federation.update_nodeinfo_period} seconds.")

    @override
    async def on_update(self) -> None:
        """Background task for polling messages to send, if no temporalio schedule was created."""
        await super().on_update()
        if self.schedule_handle:
            # if there is a temporalio schedule do not update here
            return

        params = await self.poll_delivery_queue(self.server.settings.federation)
        if not params:
            # self.logger.info(f"DeliveryWorkerTask: Polled messages, nothing to do.")
            return

        inboxes = await self.prepare_inboxes_for_delivery(params)
        results = await asyncio.gather(*[self.deliver_activity(DeliveryWorkflowParams(activity=params.activity, inbox=inbox)) for inbox in inboxes])

        self.logger.info(f"DeliveryWorkerTask: Polled messages, result status_code {DeliveryResult.create_from_list(results).status_code}.")

    @override
    async def on_shutdown(self):
        await super().on_shutdown()
        if self.schedule_handle:
            await self.schedule_handle.delete()
            self.schedule_handle = None

    @activity.defn
    async def poll_delivery_queue(self, settings_federation: FederationSettings) -> None | PrepareWorkflowParams:
        queue = self.server.delivery.delivery_queue
        if not queue:
            raise DeliveryError("Server delivery does not have a delivery queue set up to poll from")

        activity_in_queue = await queue.pop()
        if not activity_in_queue:
            return None
        activity = activity_in_queue.activity
        recipients = activity_in_queue.recipients

        return PrepareWorkflowParams(settings_federation=settings_federation, activity=activity, recipients=recipients)

    @activity.defn
    async def prepare_inboxes_for_delivery(self, params: PrepareWorkflowParams) -> list[ObjectId]:
        # deduplicate recipient list before resolving their inboxes
        recipients = list(set(params.recipients))

        # loop over recipients, carefully checking if they have an inbox
        recipient_inboxes: list[ObjectId] = []
        for recipient in recipients:
            try:
                inbox = await self.server.resolver.resolve_inboxes([recipient])
                recipient_inboxes.extend(inbox)
            except UserError as e:
                self.logger.error(f"Could not resolve inbox for recipient {recipient}: {e.message}")
            except ResolverError as e:
                self.logger.error(f"Could not resolve inbox for recipient {recipient}: {e.message}")

        # Servers MUST de-duplicate the final recipient list.
        # https://www.w3.org/TR/activitypub/#delivery
        recipient_inboxes = list(set(recipient_inboxes))

        return recipient_inboxes

    @activity.defn
    async def deliver_activity(self, params: DeliveryWorkflowParams) -> DeliveryResult:
        result = await self.server.delivery.deliver_to_url(activity=params.activity, url=params.inbox, retry_count=0)
        if result.success:
            return result
        else:
            raise ApplicationError(f"Delivery failed for inbox {params.inbox}. Result error: {result.error_message}")


@workflow.defn
class ActivityDeliveryWorkflow:

    @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


class UpdateNodeInfoTask(PeriodicBackgroundTask):

    def __init__(self, update_period: int):
        super().__init__(update_period)
        self.schedule_handle: ScheduleHandle | None = None

    @override
    async def on_startup(self) -> None:
        if not self.server:
            self.logger.warning("Trying to update node info task without initialized task.")
            return

        if not self.server.settings.worker.enabled or not self.server.worker_client:
            self.logger.warning("Trying to update node info task without enabling temporalio worker, using internal periodic task.")
            return

        self.schedule_handle = await self.server.worker_client.create_schedule(
            "workflow-schedule-UpdateNodeInfoTask",
            Schedule(
                action=ScheduleActionStartWorkflow(
                    UpdateNodeInfoWorkflow.run,
                    id="schedules-workflow-UpdateNodeInfoTask",
                    task_queue=self.server.settings.worker.tasks_queue_name,
                ),
                spec=ScheduleSpec(intervals=[ScheduleIntervalSpec(every=timedelta(seconds=self.server.settings.federation.update_nodeinfo_period))]),
                state=ScheduleState(note="Periodic update of the instance node info data."),
            ),
        )

        self.logger.info(f"Scheduled updating node info data every {self.server.settings.federation.update_nodeinfo_period} seconds.")

    @override
    async def on_update(self) -> None:
        """Background task for updating node info, if no temporalio schedule was created."""
        await super().on_update()
        if self.schedule_handle:
            # if there is a temporalio schedule, do not update here
            return

        node_info_id = await self.update_node_info()

        self.logger.info(f"UpdateNodeInfoTask: Updated node info, new id {node_info_id}")

    @override
    async def on_shutdown(self):
        await super().on_shutdown()
        if self.schedule_handle:
            await self.schedule_handle.delete()
            self.schedule_handle = None

    @activity.defn
    async def update_node_info(
        self,
    ) -> ObjectId:
        try:
            self.logger.debug("Updating node info...")
            if not self.server:
                raise DiscoveryError("Server of update info task not setup")
            settings = self.server.settings
            discovery = self.server.discovery

            collection_id = collection_id_from_name(id=settings.domain.base_url, name=UrlType.Users.value, page=None)

            node_info = await discovery.discover_nodeinfo(base_url=settings.domain.base_url)
            if node_info is None:
                raise DiscoveryError("Could not obtain NodeInfo from local instance")
            users_response = await discovery.requests.get(url=collection_id, headers=discovery.headers)
            users_response = users_response.raise_for_status()

            users_dict: dict[str, Any] = cast(dict[str, Any], users_response.json())
            users_collection = APOrderedCollection.deserialize(users_dict)
            if not isinstance(users_collection, APOrderedCollection):
                raise DiscoveryError(f"Users collection could not be retrieved properly (was type {type(users_collection)})")
            node_info.usage.users.total = users_collection.total_items or 0
            node_info.updated_at = datetime.now(timezone.utc)
            node_info.id = None # so that a new node info object gets created

            document_id = await self.server.storage.node_info.create(data=node_info)
        except ActivityPubException as e:
            raise ApplicationError(e.message)
        return document_id


@dataclass
class KeyRotationAnnounceParams:
    # The id of the new key.
    key_id_old: ObjectId
    # The id of the old key, used to sign the request for update.
    key_id_new: ObjectId
    # The actor associated to the key.
    actor_id: ObjectId


class KeyRotationTask(PeriodicBackgroundTask):

    def __init__(self, enabled: bool, update_period: int):
        super().__init__(update_period)
        self.enabled: bool = enabled
        self.schedule_handle: ScheduleHandle | None = None

    @override
    async def on_startup(self):
        await super().on_startup()

        if not self.enabled:
            self.logger.info("Key rotation task disabled.")
            return

        if not self.server:
            self.logger.warning("Trying to start key rotation task without server.")
            return

        if not self.server.settings.worker.enabled or not self.server.worker_client:
            self.logger.warning("Trying to run key rotation task without enabling temporalio worker. Schedule will not be started.")
            return

        self.schedule_handle = await self.server.worker_client.create_schedule(
            "workflow-schedule-KeyRotationTask",
            Schedule(
                action=ScheduleActionStartWorkflow(
                    KeyRotationWorkflow.run,
                    self.server.settings.federation,
                    id="schedules-workflow-KeyRotationTask",
                    task_queue=self.server.settings.worker.tasks_queue_name,
                ),
                spec=ScheduleSpec(intervals=[ScheduleIntervalSpec(every=timedelta(seconds=self.server.settings.security.key_rotation_period))]),
                state=ScheduleState(note="Periodic update of the instance key data."),
            ),
        )

        self.logger.info(f"Scheduled rotation of actor keys every {self.server.settings.security.key_rotation_period} seconds.")

    @override
    async def on_update(self) -> None:
        """Background task for key rotation, if no temporalio schedule was created."""
        if self.schedule_handle or not self.enabled:
            # if there is a temporalio schedule, or key rotation is not enabled, do not update here
            return

        await super().on_update()
        key_rotation_parameters = await self.rotate_keys()
        for params in key_rotation_parameters:
            await self.announce_key_rotation(params)
        self.logger.info(f"KeyRotationTask: Rotated instance keys.")

    @override
    async def on_shutdown(self):
        await super().on_shutdown()
        if self.server:
            await self.server.key_manager.close()
        if self.schedule_handle:
            await self.schedule_handle.delete()

    @activity.defn
    async def announce_key_rotation(self, args: KeyRotationAnnounceParams) -> None:
        """Announce new key to federation.

        This creates an APUpdate activity with the actor set to the old keys actor id,
        the new key APPublicKey object as the "object", and signs the activity proof and HTTP signature with the OLD
        key so that servers receiving it can verify that was indeed the actor that signed it.

        Args:
            args (KeyRotationAnnounceParams): Parameters for key rotation announcement.
        """
        try:
            self.logger.debug(f"Announcing key {args.key_id_new}")
            # sign the update with the old key, KEEP THE OLD KEY for a while (two months? longer?) so that the federated servers can verify the update
            key_new = await self.server.key_manager.get_public_key_APObject(key_id=args.key_id_new)

            string_to_sign = to_json_string(key_new.serialize())
            assert string_to_sign, "Could not properly serialize the new key"

            proof = await self.server.key_manager.proof_string(string_to_sign, private_key_id=args.key_id_old)

            # Announce activity, with the proof added to it.
            # The object is the new key, the actor sending it is the instance actor. It it sent to the actor on the local instance,
            # and bto will be "all federated instances", set in the deliver_to_federation method.
            announce_new_key = APAnnounce(
                actor=self.server.instance.actor_id,
                to=[args.actor_id],
                object=key_new,
                summary="Key rotation activity. Contains the new key to announce, and a proof of the serialized key object, signed with the old key.",
                proof=proof,
            )
            await self.server.delivery.deliver_to_federation(activity=announce_new_key)
            self.logger.debug(f"Completed announcing key {args.key_id_new} to federation")
        except ActivityPubException as e:
            # use the actual error message instead of user facing here, to be able to debug temporalio workflow better
            raise ApplicationError(e.message)

    @activity.defn
    async def rotate_keys(
        self,
    ) -> list[KeyRotationAnnounceParams]:
        """Checks key expiry and rotates expired keys.
        This generates a new key pair, removes expired keys, and returns parameters to announce new keys to federation in the next workflow step.

        Raises:
            KeyError: If the server of the task is not set up.

        Returns:
            list[KeyRotationAnnounceParams]: List of parameters to announce newly generated keys to the federation.
        """
        try:
            self.logger.debug("Rotating instance keys...")
            if not self.server:
                raise KeyError("Server of key rotation task not setup")

            # iterate over the pages of all actors, replace their keys if necessary
            collection_id = collection_id_from_name(id=self.server.base_url, name=ValidCollection.Users.value)
            rotation_results: list[KeyRotationAnnounceParams] = []
            async for page in self.server.collections.pages(collection_id=collection_id, access=AccessType.PRIVATE):
                # get all item ids of all keys on the page
                actor_ids = page.ordered_items or page.items or []
                actor_ids = [dereference(k, "id") for k in actor_ids if k]
                actor_ids = [k for k in actor_ids if k]

                # replace expired keys locally, asynchronously
                async def _rotate_key_of_actor(actor_id: ObjectId) -> KeyRotationAnnounceParams | None:
                    key_id = await self.server.key_manager.get_key_id_from_actor(actor_id)
                    try:
                        key = await self.server.key_manager.resolve_key(key_id=key_id)
                        if key.is_expired(datetime.now(timezone.utc)):
                            self.logger.debug(f"Rotating instance key of actor {actor_id}")
                            key_id_old = key.id
                            assert key_id_old, "key_id_old is None"
                            key_id_new = await self.server.key_manager.rotate_key(key=key, actor_id=actor_id)
                            return KeyRotationAnnounceParams(key_id_new=key_id_new, key_id_old=key_id_old, actor_id=actor_id)
                    except ActivityPubException as e:
                        self.logger.warning(f"Could not resolve key from actor {actor_id}, error {e.message}")
                    except Exception as e:
                        self.logger.warning(f"Could not resolve key from actor {actor_id}, error {e}")

                for actor_id in actor_ids:
                    try:
                        result = await _rotate_key_of_actor(actor_id)
                        if result:
                            rotation_results.append(result)
                    except ActivityPubException as e:
                        self.logger.warning(e.message)

            # TODO: remove old keys that are expired for a long time
            # TODO: make this activity "per page", so that several workers can work on different pages

            # clear key cache
            if self.server.key_manager.cache:
                self.server.key_manager.cache.clear()
            return rotation_results
        except ActivityPubException as e:
            # use the actual error message instead of user facing here, to be able to debug temporalio workflow better
            raise ApplicationError(e.message)
