import ast

import context
import asyncio

from a2a.server.events import EventQueue
from a2a.server.tasks import TaskUpdater
from a2a.types import TaskState, Part, TextPart
from a2a.utils import new_task, new_agent_text_message

from a2a_acl.agent.acl_agent import ACLAgentExecutor
from a2a_acl.content_codecs.common import natural_language_id, atom_codec_id

from a2a_acl.protocol_extension.acl_message import ACLIncomingMessage
from a2a_acl.a2a_utils.send_message import sync_reply

from a2a_acl.interface import asi_parser
from a2a_acl.protocol_extension.send_acl_message import spawn_send_acl_message
from sample_agents.llm_refactoring_advice.prompts import ask_llm_for_refactoring

my_card = asi_parser.read_file("advisor.asi").add_codecs([natural_language_id])


def decode_python(llm_answer: str):

    tag1 = "```python"
    tag2 = "```"
    if tag1 in llm_answer:
        (pre, tmp) = llm_answer.split(tag1)
        if tag2 in tmp:
            (code, post) = tmp.split(tag2)
            ast.parse(code)  # raises an exception on some syntax errors
            return (pre, code, post)
    raise ValueError(llm_answer)


class LLMAdvisor(ACLAgentExecutor):

    def __init__(self, url):
        super().__init__(agentcard=my_card, my_url=url)
        print(
            "Warning: launching a pure python A2A agent,"
            " unable to check .asi interface against python body."
        )

        self.spec = "A function that does nothing."

    async def execute_tell(
        self,
        m: ACLIncomingMessage,
        output_event_queue: EventQueue,
    ):
        self.spec = m.content
        await sync_reply(output_event_queue, "ack")
        print("Received a new specification to consider.")

    async def execute_on_upload(
        self,
        m: ACLIncomingMessage,
        output_event_queue: EventQueue,
    ):

        print("(creating a new task)")
        task = new_task(m.origin.message)
        await output_event_queue.enqueue_event(task)
        updater = TaskUpdater(output_event_queue, task.id, task.context_id)

        await asyncio.sleep(1)

        await updater.start_work()

        await asyncio.sleep(1)

        print("(consult llm)")
        res = ask_llm_for_refactoring(self.spec, m.content, context.mock_flag)
        try:
            (a, code, c) = decode_python(res)
            await updater.update_status(
                TaskState.working,
                new_agent_text_message(
                    "LLM first comment:" + a, context_id=m.origin.context_id
                ),
            )
            await updater.update_status(
                TaskState.working,
                new_agent_text_message(
                    "LLM second comment (possibly truncated):" + c,
                    context_id=m.origin.context_id,
                ),
            )

            await asyncio.sleep(1)

            await updater.add_artifact(
                [Part(root=TextPart(text=code))], name="Final response."
            )

            await asyncio.sleep(1)

            await updater.complete()

            return

        except SyntaxError:
            e = "Failed to generate syntactically correct code."
            print("ERROR: " + e)
            await self.report_error(updater, m.sender, e)

        except Exception as e:
            print("ERROR: " + str(type(e)))
            await self.report_error(updater, m.sender, "unknown reason")

    async def report_error(self, updater, to, m: str):
        # synchronous
        await updater.update_status(TaskState.failed)

        # asynchronous
        if to != "":
            spawn_send_acl_message(to, "failure", m, self.my_url, natural_language_id)
