from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.events import EventQueue
from a2a.types import AgentCard

import a2a_acl.content_codecs.common as common
from a2a_acl.protocol_extension.message_codec import acl_of_a2a
from a2a_acl.protocol_extension.acl_message import Illocution, ACLIncomingMessage
from a2a_acl.protocol_extension.send_acl_message import (
    spawn_send_acl_message,
    default_message_handler,
    default_artifact_update_handler,
)
from a2a_acl.content_codecs.common import Codec
from a2a_acl.interface.interface import ACLAgentCard

import a2a_acl.content_codecs.python_codec as python_codec
import a2a_acl.content_codecs.atom_codec as atom_codec


class ACLAgentExecutor(AgentExecutor):

    def __init__(self, agentcard: ACLAgentCard, my_url: str):
        self.public_entry_points = [
            (s.declaration_kind, s.functor) for s in agentcard.skills
        ]
        self.codecs = agentcard.supported_codecs
        self.print_entry_points()
        self.print_codecs()
        self.card = agentcard
        self.my_url = my_url
        self.codec_objects = {
            common.python_codec_id: python_codec.codec_object,
            common.atom_codec_id: atom_codec.codec_object,
        }
        self.other_message_handlers = {}

    def print_entry_points(self):
        print("Public entry points:")
        for entry_point in self.public_entry_points:
            print(" " + str(entry_point))

    def print_codecs(self):
        print("Public codecs:")
        for codec in self.codecs:
            print(" " + str(codec))

    def get_codec_object(self, codec_name: str) -> Codec:
        if codec_name in self.codec_objects:
            return self.codec_objects[codec_name]
        else:
            print("Codec " + codec_name + " not supported.")
            raise NotImplementedError

    def firewall_accept_entry_point(self, m: ACLIncomingMessage) -> bool:
        ep = (
            "*"
            if m.codec == common.natural_language_id
            else self.get_codec_object(m.codec).extract_entry_point(m.content)
        )
        if (m.illocution, ep) not in self.public_entry_points:
            print("Received a message on an entry point which is not public: " + ep)
            return False
        else:
            return True

    def firewall_accept_upload(self):
        return (Illocution.UPLOAD, "*") in self.public_entry_points

    def firewall_accept_failure(self):
        return (Illocution.FAILURE, "*") in self.public_entry_points

    def codec_accept(self, m: ACLIncomingMessage) -> bool:
        if m.codec in self.codecs:
            return True
        else:
            print("Codec not supported : " + m.codec)
            return False

    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        m: ACLIncomingMessage = acl_of_a2a(context)

        match m.illocution:

            case Illocution.TELL:
                if self.codec_accept(m) and self.firewall_accept_entry_point(m):
                    await self.execute_tell(m, event_queue)
                else:
                    print("Message not handled.")

            case Illocution.ACHIEVE:
                if self.codec_accept(m) and self.firewall_accept_entry_point(m):
                    await self.execute_achieve(m, event_queue)
                else:
                    print("Message not handled.")

            case Illocution.ASK:
                if self.codec_accept(m) and self.firewall_accept_entry_point(m):
                    await self.execute_ask(m, event_queue)
                else:
                    print("Message not handled.")

            case Illocution.UPLOAD:
                """Upload messages contain data that is not parsed by the firewall
                (no codec, no entry-point)."""
                if self.firewall_accept_upload():
                    await self.execute_on_upload(m, event_queue)
                else:
                    print("Message not handled.")

            case Illocution.FAILURE:
                """Failure messages are not parsed by the firewall
                (no codec, no entry-point)."""
                if self.firewall_accept_failure():
                    await self.execute_on_failure(m, event_queue)
                else:
                    print("Message not handled.")

            case _:
                if (
                    self.codec_accept(m)
                    and self.firewall_accept_entry_point(m)
                    and m.illocution in self.other_message_handlers
                ):
                    f = self.other_message_handlers[m.illocution]
                    await f(m, event_queue)
                else:
                    print("Unknown illocution: " + m.illocution)
                    print("Message not handled.")

    # Override-me
    async def execute_tell(
        self, m: ACLIncomingMessage, event_queue: EventQueue
    ) -> None:
        pass

    # Override-me
    async def execute_achieve(
        self, m: ACLIncomingMessage, event_queue: EventQueue
    ) -> None:
        pass

    # Override-me
    async def execute_ask(self, m: ACLIncomingMessage, event_queue: EventQueue) -> None:
        pass

    # Override-me
    async def execute_on_upload(
        self, m: ACLIncomingMessage, event_queue: EventQueue
    ) -> None:
        pass

    # Override-me
    async def execute_on_failure(
        self, m: ACLIncomingMessage, event_queue: EventQueue
    ) -> None:
        pass

    def add_message_handler(self, illoc, f):
        self.other_message_handlers[illoc] = f

    def check_public_entry_point(self, m: ACLIncomingMessage, codec: Codec):
        entry_point = codec.extract_entry_point(m.content)
        return entry_point in self.public_entry_points

    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
        raise Exception("Cancel not supported.")

    def send_message(
        self,
        dest: AgentCard | str,
        illocution: str,
        content: str,
        codec: str,
        message_processor=default_message_handler,
        task_processor=default_artifact_update_handler,
    ):
        spawn_send_acl_message(
            dest,
            illocution,
            content,
            self.my_url,
            codec,
            message_processor,
            task_processor,
        )
