import asyncio
import time
from queue import SimpleQueue

import agentspeak
import agentspeak.runtime
import agentspeak.ext_stdlib

from a2a.server.events import EventQueue
from a2a.server.tasks import TaskUpdater
from a2a.types import TaskStatusUpdateEvent, Task, TaskState, TextPart, Part
from a2a.utils import new_task
from a2a_acl import message_tools
from a2a_acl.content_codecs.common import natural_language_id
from a2a_acl.message_tools import acl_message
from agentspeak import Literal

import a2a_acl.message_tools.acl_message


from a2a_agentspeak import a2a_stdlib
from a2a_agentspeak.asp_message import build_literal, trigger, goal_type, add_message_id

from a2a_acl.agent.acl_agent import ACLAgentExecutor
from a2a_acl.message_tools.acl_message import *

from a2a_agentspeak.check import check_outgoing_illoc
from a2a_acl.message_tools.send_acl_message import (
    default_message_handler, default_artifact_update_handler,  spawn_send_acl_message,
)
from a2a_agentspeak.send_message import (
    spawn_check_recipient_intf,
)
from a2a_agentspeak.content_codecs import python_agentspeak_codec, cold_repository_codec, hot_repository_codec

from a2a_agentspeak.content_codecs.common import (
    python_agentspeak_codec_id,
    UnknownCodec,
    cold_repository_down_codec_id,
    hot_repository_down_codec_id,
    atom_codec_id,
    python_codec_id,
)
from a2a_acl.interface.interface import ACLAgentCard
from a2a_agentspeak.actions import Action

from a2a_acl.a2a_utils.send_message import sync_reply, extract_text_from_message, extract_text_from_task_artifact, build_message

def functor_of_task_state(s:TaskState)->str:
    return s.name

def final_state(t:Task)->bool | None:
    match t.status.state:
        case TaskState.completed | TaskState.failed | TaskState.canceled | TaskState.rejected:
            return True
        case TaskState.submitted | TaskState.working:
            return False
        case TaskState.input_required | TaskState.auth_required :
            return None

def send_error_handler():
    print("Error on send.")


def fast_do_send(
    to_url: str,
    illoc: Performative,
    content: agentspeak.Literal | str,
    my_url: str,
    codec: str,
    message_handler,
    artifact_handler,
    status_handler ,
    error_handler = send_error_handler # FIXME : not used.
) -> str:
    try:
        spawn_check_recipient_intf(to_url, illoc, content)
    except Exception:
        print("Error: Cannot connect to " + to_url)

    m=spawn_send_acl_message(
        dest=to_url,
        performative=str(illoc),
        content=str(content),
        sender=my_url,
        codec=codec,
        message_callback=message_handler,
        artifact_update_callback=artifact_handler,
        status_update_callback=status_handler,
        error_callback=error_handler,
    )

    return m.message_id



class EmbeddedASPAgent:
    """Encapsulates a running AgentSpeak agent."""

    def __init__(
        self,
        asl_file: str,
        url: str,
        lib: agentspeak.Actions,
        additional_actions: set[Action],
    ):
        self.queue = SimpleQueue()
        self.is_running = False
        self.current_async_tasks = []
        self.current_updaters = {}
        self.my_url = url

        self.env = agentspeak.runtime.Environment()

        # add custom actions (must occur before loading the asl file)
        self.action_lib = lib
        self.add_custom_actions()
        for t in additional_actions:
            self.add_action(t)

        with open(asl_file) as source:
            self.asp_agent = self.env.build_agent(source, self.action_lib)

    def push_call(self,a,b,c):
        self.queue.put((a,b,c))


    def pull_and_run(self):
        self.asp_agent.run()
        while not self.queue.empty():
            (a,b,c) = self.queue.get()
            self.asp_agent.call(a,b,c, agentspeak.runtime.Intention())
            self.asp_agent.run()
        d = self.asp_agent.shortest_deadline()
        if d is not None:
            time.sleep(d - time.time())
            self.pull_and_run() # Recursion


    def push_and_run(self, a, b ,c):
        self.push_call(a,b,c)
        self.run()

    def run(self):
        if not self.is_running:
            self.is_running = True
            self.pull_and_run()
            self.is_running = False



    def get_updater(self, task_id) -> TaskUpdater:
        return self.current_updaters[task_id]

    def add_updater(self, task_id, updater: TaskUpdater):
        self.current_updaters[task_id] = updater

    def remove_updater(self, task_id):
        self.current_updaters[task_id] = None

    def start(self):
        self.env.run_agent(self.asp_agent)

    # this method is called by __init__
    def add_custom_actions(self):
        actions = self.action_lib

        @actions.add_function(".send_m", (None, agentspeak.Literal, None, None))
        def _fast_send_to_url_m(
            u: agentspeak.Literal, illoc: agentspeak.Literal, term: agentspeak.Literal, callback_atom
        ) -> str:
            assert check_outgoing_illoc(illoc)

            # Use the simplest codec among atom and agentspeak
            codec = atom_codec_id if len(term.args) == 0 else python_agentspeak_codec_id

            id=fast_do_send(
                to_url=str(u),
                illoc=message_tools.acl_message.Performative(str(illoc)),
                content=term,
                my_url=self.my_url,  # Remark : here self is a ACLAgent and not an agentspeak.agent. For this reason it is difficult to get that function out of the class.
                codec=codec,
                message_handler=lambda message: self.process_message_signal(
                    extract_text_from_message(message), None
                ),
                artifact_handler=lambda task: self.process_task_artifact_update(txt=extract_text_from_task_artifact(task), callback_functor=str(callback_atom)),
                status_handler=lambda t, e, i: self.process_task_status_update(t, e, i),
                error_handler=lambda: self.process_signal(None, "send_error", "send error")
            )
            return id

        @actions.add_procedure(".send", (None, agentspeak.Literal, None))
        def _fast_send_to_url(
                u: agentspeak.Literal, illoc: agentspeak.Literal, term: agentspeak.Literal
        ):
            assert check_outgoing_illoc(illoc)

            # Use the simplest codec among atom and agentspeak
            codec = atom_codec_id if len(term.args) == 0 else python_agentspeak_codec_id

            fast_do_send(
                to_url=str(u),
                illoc=message_tools.acl_message.Performative(str(illoc)),
                content=term,
                my_url=self.my_url,
                # Remark : here self is a ACLAgent and not an agentspeak.agent. For this reason it is difficult to get that function out of the class.
                codec=codec,
                message_handler=lambda message: self.process_message_signal(
                    extract_text_from_message(message), None
                ),
                artifact_handler=lambda task: self.process_task_artifact_update(extract_text_from_task_artifact(task), None),
                status_handler=lambda t, e, i: self.process_task_status_update(t, e, i),
                error_handler=lambda: self.process_signal(None, "send_error", "send error")
            )
            return

        @actions.add_procedure(
            ".send_cb", (None, agentspeak.Literal, None, agentspeak.Literal)
        )
        def _fast_send_to_url_cb(
            u: agentspeak.Literal,
            illoc: agentspeak.Literal,
            term: agentspeak.Literal,
            cb: agentspeak.Literal,
        ):
            """Send an ACL message to an agent, specifying a callback for synchronous replies."""
            assert check_outgoing_illoc(illoc)

            # Use the simplest codec among atom and agentspeak
            codec = atom_codec_id if len(term.args) == 0 else python_agentspeak_codec_id

            fast_do_send(
                str(u),
                a2a_acl.message_tools.acl_message.Performative(str(illoc)),
                term,
                self.my_url,
                # Remark : here self is a ACLAgent and not an agentspeak.agent. For this reason it is difficult to get that function out of the class.
                codec,
                message_handler=lambda message: self.process_message_signal(
                    extract_text_from_message(message), cb.functor
                ),
                artifact_handler=lambda task: self.process_task_artifact_update(
                    extract_text_from_task_artifact(task), cb.functor
                ),
                status_handler=lambda t, e, i: self.process_task_status_update(t, e, i),
                error_handler= lambda: self.process_signal(None, "send_error", "send error")
            )

        @actions.add_procedure(".send_int", (None, agentspeak.Literal, int))
        def _fast_send_int_to_url(
            u: agentspeak.Literal, illoc: agentspeak.Literal, term: int
        ):
            assert check_outgoing_illoc(illoc)

            codec = python_codec_id

            fast_do_send(
                str(u),
                message_tools.acl_message.Performative(str(illoc)),
                str(term),
                self.my_url,
                # Remark : here self is a ACLAgent and not an agentspeak.agent. For this reason it is difficult to get that function out of the class.
                codec,
                message_handler=lambda message: self.process_message_signal(
                    extract_text_from_message(message), None
                ),
                artifact_handler=lambda task: self.process_task_artifact_update(extract_text_from_task_artifact(task), None),
                status_handler=lambda t, e, i: self.process_task_status_update(t, e, i),
                error_handler=lambda: self.process_signal(None, "send_error", "send error")
            )

        @actions.add_function(".send_bytes_cb", (None, agentspeak.Literal, str, agentspeak.Literal))
        def _fast_send_bytes_to_url_cb(
                u: agentspeak.Literal, illoc: agentspeak.Literal, term: str,
                cb: agentspeak.Literal
        ) -> str:
            assert check_outgoing_illoc(illoc)

            codec = natural_language_id

            id=fast_do_send(
                str(u),
                message_tools.acl_message.Performative(str(illoc)),
                str(term),
                self.my_url,
                # Remark : here self is a ACLAgent and not an agentspeak.agent. For this reason it is difficult to get that function out of the class.
                codec,
                message_handler=lambda message: self.process_message_signal(
                    extract_text_from_message(message), cb.functor
                ),
                artifact_handler=lambda task: self.process_task_artifact_update(extract_text_from_task_artifact(task), cb.functor),
                status_handler=lambda t, e, i: self.process_task_status_update(t, e, i),
                error_handler=lambda: self.process_signal(None, "send_error", "send error")
            )
            return id

        @actions.add_procedure(".send_bytes", (None, agentspeak.Literal, str))
        def _fast_send_bytes_to_url(
                u: agentspeak.Literal, illoc: agentspeak.Literal, term: str
        ):
            assert check_outgoing_illoc(illoc)

            codec = natural_language_id

            fast_do_send(
                str(u),
                message_tools.acl_message.Performative(str(illoc)),
                str(term),
                self.my_url,
                # Remark : here self is a ACLAgent and not an agentspeak.agent. For this reason it is difficult to get that function out of the class.
                codec,
                message_handler=default_message_handler,
                artifact_handler=default_artifact_update_handler,
                status_handler=lambda t, e, i: self.process_task_status_update(t, e, i),
                error_handler=lambda: self.process_signal(None, "send_error", "send error")
            )

        @actions.add_procedure(
            ".send_int_cb", (None, agentspeak.Literal, int, agentspeak.Literal)
        )
        def _fast_send_int_to_url_cb(
            u: agentspeak.Literal,
            illoc: agentspeak.Literal,
            term: int,
            cb: agentspeak.Literal,
        ):
            assert check_outgoing_illoc(illoc)

            codec = python_codec_id

            fast_do_send(
                str(u),
                message_tools.acl_message.Performative(str(illoc)),
                str(term),
                self.my_url,
                # Remark : here self is a ACLAgent and not an agentspeak.agent. For this reason it is difficult to get that function out of the class.
                codec,
                message_handler=lambda message: self.process_message_signal(
                    extract_text_from_message(message), cb.functor
                ),
                artifact_handler=lambda task: self.process_task_artifact_update(
                    extract_text_from_task_artifact(task), cb.functor
                ),
                status_handler=lambda t, e, i: self.process_task_status_update(t, e, i),
                error_handler=lambda: self.process_signal(None, "send_error", "send error")
            )

        @actions.add_procedure(
            ".close_task", (None)
        )
        def _close_task(
                id,
        ):
            print("Action: close task")
            t_id = str(id)
            if self.get_updater(t_id) is None :
                return
            else :
                try:
                    python_task=asyncio.create_task(self.get_updater(t_id).complete())
                    python_task.add_done_callback(self.current_async_tasks.remove)
                    self.current_async_tasks.append(python_task)
                except Exception as e:
                    print(str(e))

        @actions.add_procedure(
            ".add_artifact_on_task", (None, None)
        )
        def _add_artifact_on_task(
                id,
                term,
        ):
            print("Action: add artifact on task")
            t_id = str(id)
            if self.get_updater(t_id) is None:
                return
            else:
                try:
                    content = Part(root=TextPart(text=str(term)))
                    python_task = asyncio.create_task(self.get_updater(t_id).add_artifact(parts=[content]))
                    python_task.add_done_callback(self.current_async_tasks.remove)
                    self.current_async_tasks.append(python_task)
                except Exception as e:
                    print(str(e))

        @actions.add_procedure(
            ".reject_task", (None, None)
        )
        def _reject_task(
                id,
                term: agentspeak.Literal,
        ):
            print("Action: reject current task")
            t_id = str(id)
            if self.get_updater(t_id) is None:
                return
            else:
                try:
                    python_task=asyncio.create_task(self.get_updater(t_id).reject(build_message(str(term))))
                    python_task.add_done_callback(self.current_async_tasks.remove)
                    self.current_async_tasks.append(python_task)
                except Exception as e:
                    print(str(e))

        @actions.add_procedure(
            ".input_required", (None, None)
        )
        def _input_required(
                id,
                term: agentspeak.Literal,
        ):
            print("Action: input_required for task")
            t_id = str(id)
            if self.get_updater(t_id) is None:
                return
            else:
                try:
                    python_task = asyncio.create_task(self.get_updater(t_id).requires_input(build_message(str(term))))
                    python_task.add_done_callback(self.current_async_tasks.remove)
                    self.current_async_tasks.append(python_task)
                except Exception as e:
                    print(str(e))

        @actions.add_procedure(
            ".working_on_task", (None)
        )
        def _working_on_task(
                id,
        ):
            print("Action: work on current task")
            t_id = str(id)
            if self.get_updater(t_id) is None:
                return
            else:
                try:
                    #spawn_working_update(self.get_updater(t_id))
                    python_task = asyncio.create_task(self.get_updater(t_id).start_work())
                    python_task.add_done_callback(self.current_async_tasks.remove)
                    self.current_async_tasks.append(python_task)

                except Exception as e:
                    print(str(e))

        @actions.add_procedure(
            ".report_failure_for_task", (None, None)
        )
        def _failed_on_task(
                id, term: agentspeak.Literal
        ):
            t_id = str(id)
            print("Action: report failure for task " + t_id)

            if self.get_updater(t_id) is None:
                return
            else:
                try:
                    python_task = asyncio.create_task(self.get_updater(t_id).failed(build_message(str(term))))
                    python_task.add_done_callback(self.current_async_tasks.remove)
                    self.current_async_tasks.append(python_task)
                except Exception as e:
                    print(str(e))

    def add_action(self, action: Action):
        if action.kind == "function":
            self.action_lib.add_function(
                action.action_name, action.arity, action.implementation
            )
        else:
            print("This kind of tool is not supported yet: " + action.kind)

    async def process_tell_message(self, msg: ACLIncomingMessage):
        """Process inform requests following the AgentSpeak defined behavior."""
        assert msg.performative == TELL
        self.push_and_run(
            agentspeak.Trigger.addition,
            agentspeak.GoalType.belief,
            build_literal(msg)
        )


    async def process_achieve_message(self, msg: ACLIncomingMessage):
        """Process achieve requests following the AgentSpeak defined behavior. """
        assert msg.performative == ACHIEVE
        self.push_and_run(
            trigger(msg),
            agentspeak.GoalType.achievement,
            build_literal(msg)
        )



    def process_signal(self, txt: str | None, callback_functor: str | None, kind:str):
        if callback_functor is None:
            print("(Received synchronous " + kind + " reply : " + txt + ")")
        else:
            try:
                args = (txt,) if txt is not None else ()
                lit = Literal(callback_functor, args)
                self.push_and_run(
                    agentspeak.Trigger.addition,
                    agentspeak.GoalType.achievement,
                    lit
                )

            except agentspeak.AslError as e:
                print("Warning: " + callback_functor + " not handled (" + str(e) + ")")

    def process_task_artifact_update(self, txt: str, callback_functor: str | None):
        print("(Received a task artifact update.)")
        if callback_functor is None:
            print("(Received an artifact: " + txt + ")")
        else:
            try:
                args = (txt,) if txt is not None else ()
                lit = Literal(callback_functor, args)
                self.push_and_run(
                    agentspeak.Trigger.addition,
                    agentspeak.GoalType.achievement,
                    lit
                )
                #print("[LOG] Pushed artifact callback " +str(lit))

            except agentspeak.AslError as e:
                print("Warning: " + callback_functor + " not handled (" + str(e) + ")")

    def process_task_status_update(self, t: Task, e: TaskStatusUpdateEvent, message_id : str):
        callback_functor = functor_of_task_state(e.status.state)
        #print("([A2A]Received a task status update: " + callback_functor + ".)")

        try:
            m = t.status.message.parts[0].root.text
            args = (m, )
        except Exception :
            args = ()

        lit = add_message_id(Literal(callback_functor, args), message_id) # FIXME : add source and task_id tags
        self.push_and_run(
            agentspeak.Trigger.addition,
            agentspeak.GoalType.achievement,
            lit
        )


    def process_message_signal(self, txt: str, callback_functor: str | None):
        print("(Received a message update.)")
        self.process_signal(txt, callback_functor, "message")

    async def process_cold_repo_tell_message(self, msg: ACLIncomingMessage):
        assert msg.performative == acl_message.TELL and msg.content.startswith(
            "selected(cold_agent("
        )
        (name, holes) = cold_repository_codec.down_codec_object.decode(msg.content)
        if not holes == []:
            raise NotImplementedError()
        a = Literal("cold_agent", (name, ()))
        self.push_and_run(
            trigger(msg),
            goal_type(msg),
            Literal("selected", (a,))
        )


    async def process_hot_repo_message(self, msg: ACLIncomingMessage):
        await self.process_tell_message(msg)

    def extract_from_beliefs(self, a: str):
        r = self.asp_agent.beliefs[(a, 1)]  # fixme : arity
        assert isinstance(r, set)
        if r == set():
            return None
        else:
            tmp = next(iter(r))
            assert isinstance(tmp, agentspeak.Literal)
            assert tmp.functor == a
            assert isinstance(tmp.args, tuple)
            return tmp.args[0]

    def get_belief(self, s: str) -> str | None:
        r = self.extract_from_beliefs(s)
        if r is not None:
            return str(r)
        else:
            return None

    async def preprocess_tell_message(
        self, m: a2a_acl.message_tools.acl_message.ACLIncomingMessage, output_event_queue: EventQueue
    ):
        assert m.performative == TELL
        print("(Received an inform message: " + str(m) + ")")

        if not (m.codec == python_agentspeak_codec_id or m.codec == atom_codec_id):
            print("Warning: unknown codec " + m.codec)

        await sync_reply(output_event_queue, "Inform received.")
        await output_event_queue.close() # close the connection as soon as possible
        await self.process_tell_message(m)


    async def preprocess_achieve_message(
        self,
        m: a2a_acl.message_tools.acl_message.ACLIncomingMessage,
        output_event_queue: EventQueue
    ):
        assert m.performative == ACHIEVE
        print("(Received an achieve message: " + str(m) + ")")

        if not (m.codec == python_agentspeak_codec_id or m.codec == atom_codec_id):
            print("Warning: unknown codec " + m.codec)

        # On achieve messages, we create a new task.
        a2a_task = new_task(m.origin.message)
        python_task_1 = asyncio.create_task(output_event_queue.enqueue_event(a2a_task))
        python_task_1.add_done_callback(self.current_async_tasks.remove)
        self.current_async_tasks.append(python_task_1)
        await asyncio.sleep(0)

        updater = TaskUpdater(output_event_queue, a2a_task.id, a2a_task.context_id)

        self.add_updater(a2a_task.id, updater)

        # update to submitted
        python_task_3=asyncio.create_task(updater.submit())
        python_task_3.add_done_callback(self.current_async_tasks.remove)
        self.current_async_tasks.append(python_task_3)

        # launch behavior
        python_task_2 = asyncio.create_task(self.process_achieve_message(m))
        python_task_2.add_done_callback(self.current_async_tasks.remove)
        self.current_async_tasks.append(python_task_2)


        # don't return/close the queue before the task is completed
        # join
        while len(self.current_async_tasks) > 0 or not final_state(a2a_task):
            await asyncio.sleep(0.1)
        self.remove_updater(a2a_task.id)
        print("Request: Done.")
        return




    async def preprocess_ask_message(
        self, m: a2a_acl.message_tools.acl_message.ACLIncomingMessage, output_event_queue: EventQueue
    ):
        """Reply synchronously with the required information."""
        """in A2A, each received message has an event queue to post responses.
                This is not the case in AgentSpeak.
                Here we add a performative for requests that need an answer : ask"""
        assert m.performative == ASK  # fixme : also check that the requested belief is public.
        print("(Received an ask message: " + str(m) + ")")

        if not (m.codec == python_agentspeak_codec_id or m.codec == atom_codec_id):
            print("Warning: unknown codec " + m.codec)

        result = self.get_belief(m.content)
        if result is not None:
            await sync_reply(output_event_queue, result)
        else:
            pass  # do not reply
        await output_event_queue.close() # close only after synchronous reply is sent


    async def preprocess_cold_repo_tell_message(
        self, m: a2a_acl.message_tools.acl_message.ACLIncomingMessage, output_event_queue: EventQueue
    ):
        assert(m.performative == TELL)
        print("(Received a cold repository message: " + str(m) + ")")

        if not (m.codec == cold_repository_down_codec_id):
            print("Warning: unknown codec " + m.codec)

        await sync_reply(output_event_queue, "Tell received.")
        await self.process_cold_repo_tell_message(m)


    async def preprocess_hot_repo_tell_message(
        self, m: a2a_acl.message_tools.acl_message.ACLIncomingMessage, output_event_queue: EventQueue
    ):
        assert (m.performative == TELL)
        print("(Received a hot repository message: " + str(m) + ")")

        if not (m.codec == hot_repository_down_codec_id):
            print("Warning: unknown codec " + m.codec)

        await sync_reply(output_event_queue, "Tell received.")
        await self.process_hot_repo_message(m)



class ASPAgentExecutor(ACLAgentExecutor):

    def __init__(
        self,
        asl_file: str,
        url: str,
        additional_actions,
        card: ACLAgentCard,
    ):
        super().__init__(card, url)
        self.asl_agent = EmbeddedASPAgent(
            asl_file,
            url,
            lib=agentspeak.Actions(a2a_stdlib.actions),
            additional_actions=additional_actions,
        )
        self.codec_objects["python_agentspeak_codec"] = (
            python_agentspeak_codec.codec_object
        )
        self.codec_objects["cold_repository_down_codec"] = (
            cold_repository_codec.down_codec_object
        )

        self.codec_objects["hot_repository_down_codec"] = (
            hot_repository_codec.codec_object
        )

    def start_behavior(self):
        self.asl_agent.start()


    async def execute_tell(
        self,
        m: ACLIncomingMessage,
        output_event_queue: EventQueue,
    ) -> None:
        # FIXME (match)
        if m.codec == python_agentspeak_codec_id or m.codec == atom_codec_id:
            await self.asl_agent.preprocess_tell_message(m, output_event_queue)
        elif m.codec == cold_repository_down_codec_id:
            await self.asl_agent.preprocess_cold_repo_tell_message(
                m, output_event_queue
            )
        elif m.codec == hot_repository_down_codec_id:
            await self.asl_agent.preprocess_hot_repo_tell_message(m, output_event_queue)
        else:
            raise UnknownCodec

    async def execute_achieve(
        self,
        m: ACLIncomingMessage,
        output_event_queue: EventQueue,
    ) -> None:
        # achieve messages are always in the python_agentspeak_codec (possible misuse however because other codecs are accepted)
        await self.asl_agent.preprocess_achieve_message(m, output_event_queue)

    async def execute_ask(
        self,
        m: ACLIncomingMessage,
        output_event_queue: EventQueue,
    ) -> None:
        # ask messages are always in the python_agentspeak_codec (possible misuse however because other codecs are accepted)
        await self.asl_agent.preprocess_ask_message(m, output_event_queue)



