"""Runnable interface — shared by Agent and Team."""
from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Any, AsyncIterator
from uuid import uuid4

from logos_adk.types import MediaChunk, Response, StreamDone, TextChunk, ToolEnd, ToolStart, Usage
from logos_adk.ui_events import event_output_text, make_ui_event


class Runnable(ABC):
    """Common interface for Agent and Team."""

    @property
    @abstractmethod
    def name(self) -> str:
        ...

    @abstractmethod
    async def run(self, prompt: str, *, session_id: str | None = None, **kwargs: Any) -> Response:
        ...

    @abstractmethod
    async def stream(self, prompt: str, **kwargs: Any) -> AsyncIterator[Any]:
        ...

    async def stream_events(
        self,
        prompt: str,
        *,
        session_id: str | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[dict[str, Any]]:
        """Default UI event stream built on top of ``stream()``.

        Subclasses with richer workflow semantics should override this method.
        """
        run_id = uuid4().hex
        trace_id = kwargs.get("_trace_id")
        effective_correlation_id = correlation_id or kwargs.get("correlation_id") or session_id or run_id
        sequence = 0
        started = make_ui_event(
            event="graph_started",
            event_type="runnable",
            graph_name=self.name,
            graph_run_id=run_id,
            sequence=sequence,
            session_id=session_id,
            correlation_id=effective_correlation_id,
            trace_id=trace_id,
            node_name=self.name,
            status="running",
        )
        sequence += 1
        yield started
        yield make_ui_event(
            event="node_started",
            event_type="runnable",
            graph_name=self.name,
            graph_run_id=run_id,
            sequence=sequence,
            session_id=session_id,
            correlation_id=effective_correlation_id,
            trace_id=trace_id,
            node_name=self.name,
            status="running",
        )
        sequence += 1
        current_content = ""
        current_parts: list[dict[str, Any]] = []
        try:
            async for chunk in self.stream(
                prompt,
                session_id=session_id,
                correlation_id=effective_correlation_id,
                **kwargs,
            ):
                if isinstance(chunk, TextChunk):
                    current_content += chunk.text
                    yield make_ui_event(
                        event="node_progress",
                        event_type="text",
                        graph_name=self.name,
                        graph_run_id=run_id,
                        sequence=sequence,
                        session_id=session_id,
                        correlation_id=effective_correlation_id,
                        trace_id=trace_id,
                        node_name=self.name,
                        output=chunk.text,
                    )
                    sequence += 1
                    continue
                if isinstance(chunk, MediaChunk):
                    current_parts.append(
                        {
                            "type": chunk.part.type,
                            "text": chunk.part.text,
                            "uri": chunk.part.uri,
                            "data": chunk.part.data,
                            "mime_type": chunk.part.mime_type,
                            "name": chunk.part.name,
                            "metadata": dict(chunk.part.metadata or {}),
                        }
                    )
                    yield make_ui_event(
                        event="node_progress",
                        event_type="media",
                        graph_name=self.name,
                        graph_run_id=run_id,
                        sequence=sequence,
                        session_id=session_id,
                        correlation_id=effective_correlation_id,
                        trace_id=trace_id,
                        node_name=self.name,
                        output=current_parts[-1],
                    )
                    sequence += 1
                    continue
                if isinstance(chunk, ToolStart):
                    yield make_ui_event(
                        event="tool_started",
                        event_type="tool",
                        graph_name=self.name,
                        graph_run_id=run_id,
                        sequence=sequence,
                        session_id=session_id,
                        correlation_id=effective_correlation_id,
                        trace_id=trace_id,
                        node_name=chunk.name,
                        output=chunk.arguments,
                        metadata={"call_id": chunk.call_id},
                    )
                    sequence += 1
                    continue
                if isinstance(chunk, ToolEnd):
                    yield make_ui_event(
                        event="tool_completed",
                        event_type="tool",
                        graph_name=self.name,
                        graph_run_id=run_id,
                        sequence=sequence,
                        session_id=session_id,
                        correlation_id=effective_correlation_id,
                        trace_id=trace_id,
                        node_name=chunk.call_id,
                        status=chunk.result.status,
                        output={
                            "content": chunk.result.content,
                            "content_parts": current_parts,
                            "error": chunk.result.error,
                        },
                    )
                    sequence += 1
                    continue
                if isinstance(chunk, StreamDone):
                    response = chunk.response
                    yield make_ui_event(
                        event="node_completed",
                        event_type="runnable",
                        graph_name=self.name,
                        graph_run_id=run_id,
                        sequence=sequence,
                        session_id=session_id,
                        correlation_id=effective_correlation_id,
                        trace_id=trace_id,
                        node_name=self.name,
                        status=response.status,
                        output=response,
                    )
                    sequence += 1
                    yield make_ui_event(
                        event="graph_completed",
                        event_type="runnable",
                        graph_name=self.name,
                        graph_run_id=run_id,
                        sequence=sequence,
                        session_id=session_id,
                        correlation_id=effective_correlation_id,
                        trace_id=trace_id,
                        node_name=self.name,
                        status=response.status,
                        output=response,
                        payload_summary=event_output_text(response),
                    )
                    sequence += 1
                    return
        except Exception as exc:
            yield make_ui_event(
                event="graph_completed",
                event_type="runnable",
                graph_name=self.name,
                graph_run_id=run_id,
                sequence=sequence,
                session_id=session_id,
                correlation_id=effective_correlation_id,
                trace_id=trace_id,
                node_name=self.name,
                status="failed",
                output={"error": str(exc)},
            )
            raise
        yield make_ui_event(
            event="graph_completed",
            event_type="runnable",
            graph_name=self.name,
            graph_run_id=run_id,
            sequence=sequence,
            session_id=session_id,
            correlation_id=effective_correlation_id,
            trace_id=trace_id,
            node_name=self.name,
            status="completed",
            output=Response(
                content=current_content,
                content_parts=[],
                tool_calls=[],
                usage=Usage(input_tokens=0, output_tokens=0),
                raw={},
            ),
        )
