"""Agent wrapper that exposes a workflow through the BaseAgent surface."""

from __future__ import annotations

import json
from collections.abc import AsyncGenerator, Callable
from typing import Any, Literal

from pydantic import BaseModel

from agentbyte.agents.base import BaseAgent
from agentbyte.agents.types import AgentResponse, Usage
from agentbyte.cancellation_token import CancellationToken
from agentbyte.context import AgentContext
from agentbyte.messages import AssistantMessage, BaseMessage, Message, UserMessage
from agentbyte.middleware import MiddlewareChain
from agentbyte.workflow.core.checkpoint import (
    CheckpointConfig,
    InMemoryCheckpointStore,
    WorkflowCheckpoint,
)
from agentbyte.workflow.core.models import WorkflowExecution, WorkflowStatus
from agentbyte.workflow.core.runner import (
    WorkflowCancelledError,
    WorkflowFailedError,
    WorkflowRunner,
)

if False:  # pragma: no cover
    from agentbyte.workflow.core.workflow import Workflow


WorkflowMemoryMode = Literal["full", "last_result", "custom"]
WorkflowInputMapper = Callable[[Any, AgentContext, list[Message]], dict[str, Any]]
WorkflowOutputMapper = Callable[[WorkflowExecution, "Workflow", AgentContext], Any]
WorkflowHistorySelector = Callable[[AgentContext], list[Message]]


class WorkflowAgent(BaseAgent):
    """Expose a workflow as a BaseAgent-compatible execution surface."""

    def __init__(
        self,
        workflow: Workflow,
        *,
        name: str | None = None,
        description: str | None = None,
        memory_mode: WorkflowMemoryMode = "last_result",
        input_mapper: WorkflowInputMapper | None = None,
        output_mapper: WorkflowOutputMapper | None = None,
        history_selector: WorkflowHistorySelector | None = None,
        runner: WorkflowRunner | None = None,
    ) -> None:
        if memory_mode == "custom" and history_selector is None:
            msg = "WorkflowAgent memory_mode='custom' requires history_selector"
            raise ValueError(msg)

        self.workflow = workflow
        self.name = name or workflow.id
        self.description = (
            description
            or workflow.config.description
            or f"Workflow agent wrapper for '{workflow.config.name}'"
        )
        self.instructions = f"Workflow agent wrapper for '{workflow.config.name}'"
        self.model_client = None
        self.tools = []
        self.memory = None
        self.middleware_chain = MiddlewareChain(None)
        self.max_iterations = 1
        self.output_format = None
        self.summarize_tool_result = False
        self.required_tools = []
        self.example_tasks = list(workflow.config.metadata.get("example_tasks", []))
        self.extra_config = {}
        self._memory_mode = memory_mode
        self._input_mapper = input_mapper
        self._output_mapper = output_mapper
        self._history_selector = history_selector
        self._runner = runner or WorkflowRunner()
        self._checkpoint_cache: dict[str, WorkflowCheckpoint] = {}

    async def run(
        self,
        task: Any = None,
        context: AgentContext | None = None,
        cancellation_token: CancellationToken | None = None,
    ) -> AgentResponse:
        final_response: AgentResponse | None = None
        async for item in self.run_stream(
            task=task,
            context=context,
            cancellation_token=cancellation_token,
        ):
            if isinstance(item, AgentResponse):
                final_response = item
        if final_response is None:
            msg = "WorkflowAgent.run_stream() did not yield AgentResponse"
            raise RuntimeError(msg)
        return final_response

    async def run_stream(
        self,
        task: Any = None,
        context: AgentContext | None = None,
        cancellation_token: CancellationToken | None = None,
        verbose: bool = False,
        stream_tokens: bool = False,
    ) -> AsyncGenerator[Any, None]:
        del verbose, stream_tokens

        working_context = context if context is not None else AgentContext()
        usage = Usage()
        resume_checkpoint = self._get_resume_checkpoint(working_context)

        checkpoint_store = InMemoryCheckpointStore()
        checkpoint_config = CheckpointConfig(
            store=checkpoint_store,
            auto_save=True,
            save_interval_steps=1,
        )

        finish_reason = "stop"
        execution: WorkflowExecution | None = None

        try:
            if resume_checkpoint is not None:
                responses = self._build_resume_responses(
                    task=task,
                    pending_requests=resume_checkpoint.execution.pending_requests,
                )
                execution = await self._runner.run(
                    self.workflow,
                    cancellation_token=cancellation_token,
                    checkpoint=resume_checkpoint,
                    checkpoint_config=checkpoint_config,
                    responses=responses,
                )
            else:
                selected_history = self._select_history(working_context)
                workflow_input = self._build_workflow_input(
                    task,
                    working_context,
                    selected_history,
                )
                execution = await self._runner.run(
                    self.workflow,
                    workflow_input,
                    cancellation_token=cancellation_token,
                    checkpoint_config=checkpoint_config,
                )
            if execution.status == WorkflowStatus.SUSPENDED:
                finish_reason = "suspended"
            else:
                finish_reason = "stop"
        except WorkflowCancelledError as exc:
            finish_reason = "cancelled"
            execution = self._build_terminal_execution(WorkflowStatus.CANCELLED, str(exc))
        except WorkflowFailedError as exc:
            finish_reason = "error"
            execution = self._build_terminal_execution(WorkflowStatus.FAILED, str(exc))
        except Exception as exc:
            finish_reason = "error"
            execution = self._build_terminal_execution(WorkflowStatus.FAILED, str(exc))

        if execution is None:
            msg = "WorkflowAgent did not produce an execution result"
            raise RuntimeError(msg)

        latest_checkpoint = await checkpoint_store.load_latest(self.workflow.id)
        self._update_workflow_metadata(
            context=working_context,
            execution=execution,
            checkpoint=latest_checkpoint,
        )

        assistant_message = AssistantMessage(
            content=self._build_final_content(execution, working_context),
            source=self.name,
        )
        response_messages = self._build_response_messages(working_context, task, assistant_message)
        self._persist_messages(working_context, response_messages)
        response_context = working_context.model_copy(deep=True)
        response_context.messages = list(response_messages)

        yield assistant_message
        yield AgentResponse(
            source=self.name,
            context=response_context,
            finish_reason=finish_reason,
            usage=usage,
        )

    def _select_history(self, context: AgentContext) -> list[Message]:
        if self._memory_mode == "full":
            return list(context.messages)
        if self._memory_mode == "last_result":
            last_assistant = context.get_last_assistant_message()
            return [last_assistant] if last_assistant is not None else []
        if self._memory_mode == "custom":
            return list(self._history_selector(context))
        msg = f"Unsupported workflow agent memory mode: {self._memory_mode}"
        raise RuntimeError(msg)

    def _build_workflow_input(
        self,
        task: Any,
        context: AgentContext,
        selected_history: list[Message],
    ) -> dict[str, Any]:
        if self._input_mapper is not None:
            return self._input_mapper(task, context, selected_history)
        if isinstance(task, BaseModel):
            return task.model_dump()
        if isinstance(task, dict):
            return task.copy()

        task_text = self._normalize_task_to_text(task)
        start_step = self._get_start_step()
        fields = start_step.input_type.model_fields

        payload: dict[str, Any] = {}
        if "task" in fields:
            payload["task"] = task_text
        elif len(fields) == 1:
            only_field = next(iter(fields))
            payload[only_field] = task_text
        else:
            msg = (
                f"WorkflowAgent '{self.name}' cannot derive workflow input for "
                f"start step '{start_step.step_id}'. Provide input_mapper."
            )
            raise RuntimeError(msg)

        if "history" in fields:
            payload["history"] = [message.model_dump(mode="json") for message in selected_history]
        if "messages" in fields:
            payload["messages"] = [message.model_dump(mode="json") for message in selected_history]

        return payload

    def _build_resume_responses(
        self,
        *,
        task: Any,
        pending_requests: dict[str, Any],
    ) -> dict[str, Any]:
        if not pending_requests:
            msg = f"WorkflowAgent '{self.name}' has no pending requests to resume"
            raise RuntimeError(msg)

        if isinstance(task, dict):
            if set(task).issubset(set(pending_requests)):
                return task.copy()
            if len(pending_requests) == 1:
                request_id = next(iter(pending_requests))
                return {request_id: task.copy()}

        if isinstance(task, BaseModel):
            if len(pending_requests) == 1:
                request_id = next(iter(pending_requests))
                return {request_id: task.model_dump()}

        if len(pending_requests) == 1:
            request_id = next(iter(pending_requests))
            return {request_id: self._normalize_task_to_text(task)}

        msg = (
            f"WorkflowAgent '{self.name}' requires explicit responses keyed by request id "
            "when resuming multiple pending workflow requests"
        )
        raise RuntimeError(msg)

    @staticmethod
    def _normalize_task_to_text(task: Any) -> str:
        if task is None:
            return ""
        if isinstance(task, str):
            return task
        if isinstance(task, UserMessage):
            return task.content
        if isinstance(task, list):
            parts = [message.content for message in task if hasattr(message, "content")]
            return "\n".join(parts)
        return str(task)

    def _build_final_content(
        self,
        execution: WorkflowExecution,
        context: AgentContext,
    ) -> str:
        if execution.status == WorkflowStatus.SUSPENDED:
            return "Workflow suspended awaiting external input."
        if execution.status == WorkflowStatus.FAILED:
            return f"Workflow execution failed: {execution.error or 'unknown error'}"
        if execution.status == WorkflowStatus.CANCELLED:
            return f"Workflow execution cancelled: {execution.error or 'cancelled'}"

        if self._output_mapper is not None:
            output_value = self._output_mapper(execution, self.workflow, context)
        else:
            output_value = self._extract_default_output(execution)
        return self._stringify_output(output_value)

    def _extract_default_output(self, execution: WorkflowExecution) -> Any:
        if len(self.workflow.end_step_ids) != 1:
            msg = (
                f"WorkflowAgent '{self.name}' requires exactly one workflow end step "
                "for default output extraction; provide output_mapper for custom behavior"
            )
            raise RuntimeError(msg)
        end_step_id = self.workflow.end_step_ids[0]
        step_execution = execution.step_executions.get(end_step_id)
        if step_execution is None or step_execution.output_data is None:
            msg = (
                f"WorkflowAgent '{self.name}' could not read output from end step "
                f"'{end_step_id}'"
            )
            raise RuntimeError(msg)
        return step_execution.output_data

    @staticmethod
    def _stringify_output(value: Any) -> str:
        if isinstance(value, AssistantMessage):
            return value.content
        if isinstance(value, str):
            return value
        if isinstance(value, BaseModel):
            value = value.model_dump(mode="json")
        if isinstance(value, dict):
            for key in ("response", "result", "message", "content"):
                field_value = value.get(key)
                if isinstance(field_value, str):
                    return field_value
            return json.dumps(value, sort_keys=True)
        return str(value)

    @staticmethod
    def _build_response_messages(
        context: AgentContext,
        task: Any,
        assistant_message: AssistantMessage,
    ) -> list[Message]:
        combined = list(context.messages)
        combined.extend(WorkflowAgent._task_to_messages(task))
        combined.append(assistant_message)
        return combined

    def _persist_messages(
        self,
        context: AgentContext,
        combined_messages: list[Message],
    ) -> None:
        temp_context = context.model_copy(deep=True)
        temp_context.messages = list(combined_messages)
        context.messages = self._select_history(temp_context)

    @staticmethod
    def _task_to_messages(task: Any) -> list[Message]:
        if task is None:
            return []
        if isinstance(task, str):
            return [UserMessage(content=task, source="user")]
        if isinstance(task, UserMessage):
            return [task]
        if isinstance(task, list):
            return [message for message in task if isinstance(message, BaseMessage)]
        return []

    def _update_workflow_metadata(
        self,
        *,
        context: AgentContext,
        execution: WorkflowExecution,
        checkpoint: Any,
    ) -> None:
        previous_checkpoint_id = (
            context.metadata.get("workflow_agent", {}).get("checkpoint_id")
            if isinstance(context.metadata.get("workflow_agent"), dict)
            else None
        )
        checkpoint_id: str | None = None
        if execution.status == WorkflowStatus.SUSPENDED and checkpoint is not None:
            checkpoint_id = checkpoint.checkpoint_id
            self._checkpoint_cache[checkpoint_id] = checkpoint
        elif previous_checkpoint_id is not None:
            self._checkpoint_cache.pop(previous_checkpoint_id, None)

        workflow_meta = {
            "workflow_id": self.workflow.id,
            "execution_id": execution.id,
            "workflow_status": execution.status.value,
            "checkpoint_id": checkpoint_id,
            "pending_requests": {
                request_id: request.model_dump(mode="json")
                for request_id, request in execution.pending_requests.items()
            },
        }
        context.metadata["workflow_agent"] = workflow_meta

    def _get_resume_checkpoint(self, context: AgentContext) -> WorkflowCheckpoint | None:
        workflow_meta = context.metadata.get("workflow_agent")
        if not isinstance(workflow_meta, dict):
            return None
        if workflow_meta.get("workflow_id") != self.workflow.id:
            return None
        if workflow_meta.get("workflow_status") != WorkflowStatus.SUSPENDED.value:
            return None
        checkpoint_id = workflow_meta.get("checkpoint_id")
        if not checkpoint_id:
            return None

        checkpoint = self._checkpoint_cache.get(checkpoint_id)
        if checkpoint is None:
            msg = (
                f"WorkflowAgent '{self.name}' cannot resume suspended workflow "
                f"because checkpoint '{checkpoint_id}' is not available in memory"
            )
            raise RuntimeError(msg)
        return checkpoint

    def _get_start_step(self) -> Any:
        if self.workflow.start_step_id is None:
            msg = f"WorkflowAgent '{self.name}' cannot run a workflow without a start step"
            raise RuntimeError(msg)
        return self.workflow.steps[self.workflow.start_step_id]

    def _build_terminal_execution(
        self,
        status: WorkflowStatus,
        error: str,
    ) -> WorkflowExecution:
        return WorkflowExecution(
            workflow_id=self.workflow.id,
            status=status,
            error=error,
        )


__all__ = ["WorkflowAgent", "WorkflowMemoryMode"]
