from asyncio import CancelledError, Task
from collections.abc import AsyncIterator, Sequence
from typing import Any
from uuid import UUID, uuid4

from haiway import AsyncQueue, ctx

from draive.conversation.types import (
    ConversationEvent,
    ConversationMessage,
    ConversationMessageChunk,
    ConversationStreamElement,
    RealtimeConversationMemory,
)
from draive.instructions import Instruction
from draive.lmm import (
    LMMSession,
    LMMSessionEvent,
    LMMSessionOutput,
    LMMStreamChunk,
    LMMStreamInput,
    LMMToolRequest,
    LMMToolResponse,
)
from draive.lmm.types import LMMContextElement
from draive.multimodal import MetaContent
from draive.tools import Toolbox
from draive.utils import Processing, ProcessingEvent

__all__ = ("realtime_conversation",)


async def realtime_conversation(  # noqa: C901, PLR0915
    *,
    instruction: Instruction | None,
    input_stream: AsyncIterator[ConversationMessageChunk],
    memory: RealtimeConversationMemory,
    toolbox: Toolbox,
    **extra: Any,
) -> AsyncIterator[ConversationStreamElement]:
    initial_context: Sequence[LMMContextElement] = tuple(
        message.as_lmm_context_element()
        for message in await memory.recall()  # relying on memory recall correctness
        if isinstance(message, ConversationMessage)  # skip events in LMM context
    )

    input_queue = AsyncQueue[LMMStreamInput]()

    async def process_input() -> None:
        try:
            ctx.log_debug("Starting input processing...")
            # TODO: include guardrails check
            async for element in input_stream:
                if isinstance(element, ConversationEvent):
                    continue  # skip events in LMM context

                input_queue.enqueue(
                    LMMStreamChunk.of(
                        element.content,
                        eod=element.eod,
                    )
                )
                await memory.remember(element)  # add to memory

            input_queue.finish()
            ctx.log_debug("...input processing finished!")

        except CancelledError:
            # just finish input on cancel
            input_queue.finish()

        except BaseException as exc:
            input_queue.finish(exception=exc)

    pending_tool_requests: dict[str, Task[None]] = {}

    async def handle_tool_request(
        tool_request: LMMToolRequest,
        /,
    ) -> None:
        try:
            ctx.log_debug(f"Handling tool request ({tool_request.identifier})...")
            response: LMMToolResponse = await toolbox.respond(tool_request)
            # check if task was not cancelled before passing the result to input
            ctx.check_cancellation()
            # deliver the result directly to input
            input_queue.enqueue(response)

        except CancelledError:
            pass  # just cancelled, ignore

        finally:
            # update pending tools - we have delivered or ignored the result
            del pending_tool_requests[tool_request.identifier]
            ctx.log_debug(f"...tool request ({tool_request.identifier}) handling finished!")

    input_processing_task: Task[None] = ctx.spawn(process_input)

    session_output: AsyncIterator[LMMSessionOutput]
    try:
        ctx.log_debug("Preparing LMM session...")
        session_output = await LMMSession.prepare(
            instruction=instruction,
            input_stream=input_queue,
            initial_context=initial_context,
            tools=toolbox.available_tools(),
        )
        del initial_context

    except BaseException as exc:
        ctx.log_error(
            "...failed to prepare LMM session!",
            exception=exc,
        )
        input_processing_task.cancel()
        raise exc

    output_queue = AsyncQueue[ConversationStreamElement]()
    # current output message identifier to match chunks and events
    output_message_identifier: UUID = uuid4()

    async def report_event(
        event: ProcessingEvent,
        /,
    ) -> None:
        ctx.log_debug("Processing conversation event")
        output_queue.enqueue(
            ConversationEvent.of(
                event.name,
                message_identifier=output_message_identifier,
                identifier=event.identifier,
                content=event.content,
                meta=event.meta,
            )
        )

    with ctx.updated(ctx.state(Processing).updated(event_reporting=report_event)):

        async def process_output() -> None:
            nonlocal output_message_identifier
            try:
                ctx.log_debug("Starting output processing...")
                async for chunk in session_output:
                    match chunk:
                        case LMMStreamChunk() as chunk:
                            output_chunk: ConversationMessageChunk = ConversationMessageChunk.model(
                                chunk.content,
                                message_identifier=output_message_identifier,
                                eod=chunk.eod,
                            )
                            # send to output
                            output_queue.enqueue(output_chunk)
                            if chunk.eod:
                                # start next message when current finished
                                output_message_identifier = uuid4()
                            # and add to memory
                            await memory.remember(output_chunk)

                        case LMMSessionEvent() as event:
                            conversation_event: ConversationEvent
                            match event.category:
                                case "interrupted":
                                    # cancel pending tools on interrupt
                                    for task in pending_tool_requests.values():
                                        task.cancel()

                                    conversation_event = ConversationEvent.of(
                                        "interrupted",
                                        message_identifier=output_message_identifier,
                                        meta=event.meta,
                                    )

                                case "completed":
                                    conversation_event = ConversationEvent.of(
                                        "completed",
                                        message_identifier=output_message_identifier,
                                        meta=event.meta,
                                    )

                                case other:
                                    conversation_event = ConversationEvent.of(
                                        other,
                                        message_identifier=output_message_identifier,
                                        content=MetaContent.of(other),
                                        meta=event.meta,
                                    )

                            # send to output
                            output_queue.enqueue(conversation_event)
                            # and add to memory
                            await memory.remember(conversation_event)

                        case LMMToolRequest() as tool_request:
                            pending_tool_requests[tool_request.identifier] = ctx.spawn(
                                handle_tool_request,
                                tool_request,
                            )

            finally:
                # cancel input reading
                input_processing_task.cancel()
                # cancel pending tools
                for task in pending_tool_requests.values():
                    task.cancel()

                ctx.log_debug("...output processing finished!")

        ctx.spawn(process_output)

    return output_queue
