import asyncio
import uuid
from threading import Thread
from typing import Callable, Any

import a2a
import httpx
from a2a.client import ClientFactory, ClientConfig
from a2a.server.events import EventQueue
from a2a.types import (
    Message,
    Role,
    Part,
    TextPart,
    AgentCard,
    Task,
    PushNotificationConfig,
    TaskStatusUpdateEvent,
)

from a2a.utils import new_agent_text_message

from a2a_acl.a2a_utils.card_holder import download_card


def build_message(content: str) -> Message:
    return Message(
        role=Role.agent,
        message_id=uuid.uuid4().hex,
        parts=[Part(root=TextPart(text=content))],
    )


class SendFailureException(Exception):
    pass


def extract_text_from_message(m: Message) -> str:
    return m.parts[0].root.text


def extract_text_from_task_artifact(t: Task) -> str:
    res = str(t.status.state)
    if t.artifacts:
        a = t.artifacts[-1]
        for p in a.parts:
            if p.root.kind == "text":
                res = p.root.text
    return res


def default_message_handler(m: Message) -> None:
    print("(Message event received: " + extract_text_from_message(m) + ")")


def default_artifact_update_handler(t: Task) -> None:
    print("(Artifact event received: " + extract_text_from_task_artifact(t) + ")")


def default_status_update_handler(
    t: Task, e: TaskStatusUpdateEvent, message_id: str
) -> None:
    print(
        "(Status update event received for task "
        + t.id
        + " initiated by message "
        + message_id
        + " : "
        + str(e.status.state)
        + ")"
    )


def default_error_callback():
    print(
        "Error while talking with agent (you can specify a callback to handle this error)."
    )


# derived from https://github.com/a2aproject/a2a-samples/blob/main/samples/python/agents/number_guessing_game/utils/protocol_wrappers.py
async def send_message(
    destination: AgentCard | str,
    msg: Message,
    sender: str,
    message_callback: Callable[[Message], None] = default_message_handler,
    artifact_update_callback: Callable[[Task], None] = default_artifact_update_handler,
    status_update_callback: Callable[
        [Any, a2a.types.TaskStatusUpdateEvent, str], None
    ] = default_status_update_handler,
):
    if isinstance(destination, str):
        try:
            destination = await download_card(destination)
        except Exception:
            print("Error: Cannot get agent card. Send failed.")
            raise SendFailureException

    else:
        assert isinstance(destination, AgentCard)

    _client_factory = ClientFactory(
        ClientConfig(
            push_notification_configs=[PushNotificationConfig(url=sender)],
            httpx_client=httpx.AsyncClient(timeout=httpx.Timeout(timeout=30)),
            streaming=destination.capabilities.streaming,
        )
    )
    client = _client_factory.create(destination)

    try:
        # Warning : concurrent async generators work badly in python
        # -> two concurrent message sending should be in two separate threads
        # (see spawn_send_message below)
        async for event in client.send_message(msg):
            if isinstance(event, Message):
                message_callback(event)

            else:
                assert isinstance(event, tuple)
                (t, opt) = event
                assert isinstance(t, Task)
                match type(opt):
                    case a2a.types.TaskStatusUpdateEvent:
                        status_update_callback(t, opt, msg.message_id)
                    case a2a.types.TaskArtifactUpdateEvent:
                        artifact_update_callback(t)
                    case None:
                        artifact_update_callback(t)
        print("End of sending")

    except a2a.client.errors.A2AClientTimeoutError:
        print(
            "Warning: no synchronous reply before timeout. Some information might be lost."
        )
    except a2a.client.errors.A2AClientHTTPError as e:
        print("Send failed (" + e.message + ")")
        raise SendFailureException

    except Exception as e:
        print("Send failed. (" + str(type(e)) + ", " + str(e) + ")")
        raise SendFailureException


# hard reference to task that must not be garbage collected.
launched_tasks = []
# Note that if we never delete tasks from this list, we create a memory leak. FIXME


class SendMessageThread(Thread):

    def __init__(
        self,
        dest: AgentCard | str,
        msg: Message,
        sender: str,
        error_callback,
        message_callback=default_message_handler,
        artifact_update_callback=default_artifact_update_handler,
        status_update_callback=default_status_update_handler,
    ):
        super().__init__()
        self.dest = dest
        self.msg = msg
        self.reply_to = sender
        self.message_callback = message_callback
        self.artifact_update_callback = artifact_update_callback
        self.status_update_callback = status_update_callback
        self.error_callback = error_callback

    def run(self):
        try:
            loop = asyncio.new_event_loop()
            t = loop.create_task(
                send_message(
                    self.dest,
                    self.msg,
                    self.reply_to,
                    message_callback=self.message_callback,
                    artifact_update_callback=self.artifact_update_callback,
                    status_update_callback=self.status_update_callback,
                )
            )
            launched_tasks.append(t)
            # TRY-ME : t.add_done_callback(launched_tasks.remove)
            loop.run_until_complete(t)
            print("(end of send thread)")

        except SendFailureException:
            print("Send failed.")
            self.error_callback()


def spawn_send_message(
    dest: AgentCard | str,
    content: str,
    sender: str,
    message_callback=default_message_handler,
    task_callback=default_artifact_update_handler,
    status_update_callback=default_status_update_handler,
    error_callback=default_error_callback,
) -> None:
    msg = build_message(content)
    t = SendMessageThread(
        dest,
        msg,
        sender,
        error_callback,
        message_callback,
        task_callback,
        status_update_callback,
    )
    t.start()


async def sync_reply(output_event_queue: EventQueue, r: str):
    # Warning : according to A2A spec v1.0rc1, you can reply with only 1 Message then close, or a stream of Task.
    await output_event_queue.enqueue_event(new_agent_text_message(r))
