"""Execution engine for Agentbyte WebUI."""

from __future__ import annotations

from collections.abc import AsyncGenerator
from datetime import datetime
from typing import Any

from pydantic import BaseModel, TypeAdapter

from agentbyte.agents.base import BaseAgent
from agentbyte.agents.types import AgentResponse
from agentbyte.cancellation_token import CancellationToken
from agentbyte.context import AgentContext, ToolApprovalResponse
from agentbyte.messages import Message
from agentbyte.orchestration.base import BaseOrchestrator
from agentbyte.types import OrchestrationResponse
from agentbyte.workflow.core.models import WorkflowExecution
from agentbyte.workflow.core.runner import WorkflowRunner
from agentbyte.workflow.core.workflow import BaseWorkflow

from .models import WebUIStreamEvent
from .sessions import SessionBindingError, SessionManager

MESSAGE_LIST_ADAPTER = TypeAdapter(list[Message])


def _serialize_payload(payload: Any) -> Any:
    """Serialize framework payloads for API responses and SSE."""
    if isinstance(payload, BaseModel):
        return payload.model_dump(mode="json")
    if isinstance(payload, datetime):
        return payload.isoformat()
    if isinstance(payload, list):
        return [_serialize_payload(item) for item in payload]
    if isinstance(payload, tuple):
        return [_serialize_payload(item) for item in payload]
    if isinstance(payload, dict):
        return {str(key): _serialize_payload(value) for key, value in payload.items()}
    return payload


def _sse(event: WebUIStreamEvent) -> str:
    return f"data: {event.model_dump_json()}\n\n"


class ExecutionEngine:
    """Execute Agentbyte entities with session continuity."""

    def __init__(self, session_manager: SessionManager) -> None:
        self.session_manager = session_manager

    def parse_messages(self, messages: list[dict[str, Any]] | None) -> list[Message]:
        """Parse serialized messages into Agentbyte message objects."""
        if not messages:
            return []
        return MESSAGE_LIST_ADAPTER.validate_python(messages)

    async def execute_agent(
        self,
        entity_id: str,
        agent: BaseAgent,
        messages: list[Message],
        session_id: str | None = None,
        user_id: str | None = None,
        approval_responses: list[ToolApprovalResponse] | None = None,
        cancellation_token: CancellationToken | None = None,
    ) -> AgentResponse:
        """Run an agent to completion."""
        session_id = session_id or self.session_manager.create_session_id()
        context = await self.session_manager.get_or_create(
            session_id,
            entity_id,
            "agent",
            user_id=user_id,
        )
        self._append_messages(context, messages)
        self._inject_approvals(context, approval_responses)
        response = await agent.run(
            task=None,
            context=context,
            cancellation_token=cancellation_token,
        )
        await self.session_manager.update(session_id, response.context)
        return response

    async def execute_agent_stream(
        self,
        entity_id: str,
        agent: BaseAgent,
        messages: list[Message],
        session_id: str | None = None,
        user_id: str | None = None,
        stream_tokens: bool = True,
        approval_responses: list[ToolApprovalResponse] | None = None,
        cancellation_token: CancellationToken | None = None,
    ) -> AsyncGenerator[str, None]:
        """Stream agent execution as SSE."""
        session_id = session_id or self.session_manager.create_session_id()
        context = await self.session_manager.get_or_create(
            session_id,
            entity_id,
            "agent",
            user_id=user_id,
        )
        self._append_messages(context, messages)
        self._inject_approvals(context, approval_responses)

        try:
            async for item in agent.run_stream(
                task=None,
                context=context,
                cancellation_token=cancellation_token,
                verbose=True,
                stream_tokens=stream_tokens,
            ):
                yield _sse(
                    WebUIStreamEvent(
                        session_id=session_id,
                        event=_serialize_payload(item),
                    )
                )
            await self.session_manager.update(session_id, context)
        except SessionBindingError:
            raise
        except Exception as exc:
            yield _sse(
                WebUIStreamEvent(
                    session_id=session_id,
                    event={"type": "error", "message": str(exc)},
                )
            )

    async def execute_orchestrator(
        self,
        entity_id: str,
        orchestrator: BaseOrchestrator,
        messages: list[Message],
        session_id: str | None = None,
        user_id: str | None = None,
        cancellation_token: CancellationToken | None = None,
    ) -> OrchestrationResponse:
        """Run an orchestrator to completion."""
        session_id = session_id or self.session_manager.create_session_id()
        context = await self.session_manager.get_or_create(
            session_id,
            entity_id,
            "orchestrator",
            user_id=user_id,
        )
        # Pass new messages as task and prior session history via context.
        # The orchestrator seeds shared_messages from context, appends the task,
        # runs the loop, then writes the full conversation back to context.
        response = await orchestrator.run(
            task=messages,
            cancellation_token=cancellation_token,
            context=context,
        )
        context.metadata["final_result"] = response.final_result
        context.metadata["stop_reason"] = response.stop_message.content
        await self.session_manager.update(session_id, context)
        return response

    async def execute_orchestrator_stream(
        self,
        entity_id: str,
        orchestrator: BaseOrchestrator,
        messages: list[Message],
        session_id: str | None = None,
        user_id: str | None = None,
        cancellation_token: CancellationToken | None = None,
    ) -> AsyncGenerator[str, None]:
        """Stream orchestrator execution as SSE."""
        session_id = session_id or self.session_manager.create_session_id()
        context = await self.session_manager.get_or_create(
            session_id,
            entity_id,
            "orchestrator",
            user_id=user_id,
        )
        # Pass new messages as task and prior session history via context.
        # The orchestrator updates context.messages in-place when done.
        try:
            async for item in orchestrator.run_stream(
                task=messages,
                cancellation_token=cancellation_token,
                verbose=True,
                context=context,
            ):
                if isinstance(item, OrchestrationResponse):
                    context.metadata["final_result"] = item.final_result
                    context.metadata["stop_reason"] = item.stop_message.content
                    await self.session_manager.update(session_id, context)
                yield _sse(
                    WebUIStreamEvent(
                        session_id=session_id,
                        event=_serialize_payload(item),
                    )
                )
        except SessionBindingError:
            raise
        except Exception as exc:
            yield _sse(
                WebUIStreamEvent(
                    session_id=session_id,
                    event={"type": "error", "message": str(exc)},
                )
            )

    async def execute_workflow(
        self,
        entity_id: str,
        workflow: BaseWorkflow,
        input_data: Any,
        session_id: str | None = None,
        user_id: str | None = None,
        cancellation_token: CancellationToken | None = None,
    ) -> WorkflowExecution:
        """Run a workflow to completion."""
        session_id = session_id or self.session_manager.create_session_id()
        context = await self.session_manager.get_or_create(
            session_id,
            entity_id,
            "workflow",
            user_id=user_id,
        )
        context.metadata["last_input"] = _serialize_payload(input_data)
        runner = WorkflowRunner()
        execution = await runner.run(
            workflow=workflow,
            initial_input=input_data,
            cancellation_token=cancellation_token,
        )
        context.metadata["last_execution"] = execution.model_dump(mode="json")
        await self.session_manager.update(session_id, context)
        return execution

    async def execute_workflow_stream(
        self,
        entity_id: str,
        workflow: BaseWorkflow,
        input_data: Any,
        session_id: str | None = None,
        user_id: str | None = None,
        cancellation_token: CancellationToken | None = None,
    ) -> AsyncGenerator[str, None]:
        """Stream workflow execution as SSE."""
        session_id = session_id or self.session_manager.create_session_id()
        context = await self.session_manager.get_or_create(
            session_id,
            entity_id,
            "workflow",
            user_id=user_id,
        )
        context.metadata["last_input"] = _serialize_payload(input_data)
        runner = WorkflowRunner()
        try:
            async for item in runner.run_stream(
                workflow=workflow,
                initial_input=input_data,
                cancellation_token=cancellation_token,
            ):
                if hasattr(item, "execution") and getattr(item, "execution", None) is not None:
                    context.metadata["last_execution"] = _serialize_payload(item.execution)
                    await self.session_manager.update(session_id, context)
                yield _sse(
                    WebUIStreamEvent(
                        session_id=session_id,
                        event=_serialize_payload(item),
                    )
                )
        except SessionBindingError:
            raise
        except Exception as exc:
            yield _sse(
                WebUIStreamEvent(
                    session_id=session_id,
                    event={"type": "error", "message": str(exc)},
                )
            )

    def _append_messages(self, context: AgentContext, messages: list[Message]) -> None:
        if not messages:
            return
        if context.messages:
            last = context.messages[-len(messages):]
            if _serialize_payload(last) == _serialize_payload(messages):
                return
        for message in messages:
            context.add_message(message)

    def _inject_approvals(
        self,
        context: AgentContext,
        approval_responses: list[ToolApprovalResponse] | None,
    ) -> None:
        for response in approval_responses or []:
            context.add_approval_response(response)


__all__ = ["ExecutionEngine"]
