"""SubWorkflowStep executes a child workflow inside a parent workflow step."""

from __future__ import annotations

import copy
from typing import TYPE_CHECKING, Any, Callable
from uuid import uuid4

from pydantic import BaseModel

from agentbyte.workflow.core.checkpoint import (
    CheckpointConfig,
    InMemoryCheckpointStore,
    WorkflowCheckpoint,
)
from agentbyte.workflow.core.models import (
    WorkflowContext,
    WorkflowExecution,
    WorkflowPendingRequest,
    WorkflowRequestKind,
    _WorkflowRequestSuspended,
)
from agentbyte.workflow.steps.step import BaseStep

if TYPE_CHECKING:
    from agentbyte.workflow.core.workflow import Workflow


ChildInputMapper = Callable[[Any, WorkflowContext], dict[str, Any]]
ChildOutputMapper = Callable[[WorkflowExecution, "Workflow", WorkflowContext], Any]
WorkflowBuilder = Any


class SubWorkflowStep(BaseStep[Any, Any]):
    """Run a child workflow as a regular workflow step."""

    def __init__(
        self,
        step_id: str,
        metadata,
        workflow: WorkflowBuilder,
        input_type,
        output_type,
        *,
        input_mapper: ChildInputMapper | None = None,
        output_mapper: ChildOutputMapper | None = None,
        child_max_concurrent_steps: int = 5,
    ) -> None:
        super().__init__(step_id, metadata, input_type, output_type)
        self._workflow_builder = workflow
        self._input_mapper = input_mapper
        self._output_mapper = output_mapper
        self._child_max_concurrent_steps = child_max_concurrent_steps
        self._active_invocations: dict[str, WorkflowCheckpoint] = {}
        self._request_routes: dict[str, dict[str, str]] = {}

    async def execute(self, input_data: Any, context: WorkflowContext) -> Any:
        self._active_invocations.clear()
        self._request_routes.clear()
        invocation_id = str(uuid4())
        child_workflow = self._create_child_workflow()
        child_execution, child_checkpoint = await self._run_child_workflow(
            child_workflow=child_workflow,
            initial_input=self._map_child_input(input_data, context),
        )
        if child_execution.status.value == "suspended":
            if child_checkpoint is None:
                msg = "Sub-workflow suspended without producing a checkpoint"
                raise RuntimeError(msg)
            self._active_invocations[invocation_id] = child_checkpoint
            await self._propagate_child_request(
                context=context,
                invocation_id=invocation_id,
                child_checkpoint=child_checkpoint,
                child_execution=child_execution,
            )
        return self._coerce_output(
            self._extract_child_output(
                execution=child_execution,
                child_workflow=child_workflow,
                context=context,
            )
        )

    async def resume_from_request(
        self,
        input_data: Any,
        context: WorkflowContext,
        request: WorkflowPendingRequest,
        response: Any,
    ) -> Any:
        route = self._request_routes.pop(request.request_id, None)
        if route is None:
            msg = f"Missing sub-workflow route for request '{request.request_id}'"
            raise RuntimeError(msg)

        invocation_id = route["invocation_id"]
        child_request_id = route["child_request_id"]
        child_checkpoint = self._active_invocations.get(invocation_id)
        if child_checkpoint is None:
            msg = f"Missing child checkpoint for sub-workflow invocation '{invocation_id}'"
            raise RuntimeError(msg)

        child_workflow = self._create_child_workflow()
        child_execution, latest_child_checkpoint = await self._run_child_workflow(
            child_workflow=child_workflow,
            checkpoint=child_checkpoint,
            responses={child_request_id: response},
        )
        if child_execution.status.value == "suspended":
            if latest_child_checkpoint is None:
                msg = "Sub-workflow suspended without producing a checkpoint"
                raise RuntimeError(msg)
            self._active_invocations[invocation_id] = latest_child_checkpoint
            await self._propagate_child_request(
                context=context,
                invocation_id=invocation_id,
                child_checkpoint=latest_child_checkpoint,
                child_execution=child_execution,
            )

        self._active_invocations.pop(invocation_id, None)
        return self._coerce_output(
            self._extract_child_output(
                execution=child_execution,
                child_workflow=child_workflow,
                context=context,
            )
        )

    def on_checkpoint_save(self) -> dict[str, Any]:
        if not self._active_invocations and not self._request_routes:
            return {}
        return {
            "active_invocations": {
                invocation_id: checkpoint.model_dump(mode="json")
                for invocation_id, checkpoint in self._active_invocations.items()
            },
            "request_routes": copy.deepcopy(self._request_routes),
        }

    def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
        self._active_invocations = {
            invocation_id: WorkflowCheckpoint.model_validate(checkpoint_data)
            for invocation_id, checkpoint_data in state.get("active_invocations", {}).items()
        }
        self._request_routes = {
            request_id: dict(route)
            for request_id, route in state.get("request_routes", {}).items()
        }

    def _create_child_workflow(self) -> Workflow:
        from agentbyte.workflow.core.workflow import Workflow

        builder = self._workflow_builder
        if callable(builder):
            workflow = builder()
        else:
            try:
                workflow = copy.deepcopy(builder)
            except Exception as exc:
                msg = (
                    "Failed to clone child workflow for SubWorkflowStep. "
                    "Pass a workflow factory instead of a live workflow instance "
                    "when the child graph contains non-copyable objects."
                )
                raise RuntimeError(msg) from exc
        if not isinstance(workflow, Workflow):
            msg = f"SubWorkflowStep expected Workflow, got {type(workflow)!r}"
            raise RuntimeError(msg)
        return workflow

    def _map_child_input(self, input_data: Any, context: WorkflowContext) -> dict[str, Any]:
        if self._input_mapper is not None:
            return self._input_mapper(input_data, context)
        if isinstance(input_data, BaseModel):
            return input_data.model_dump()
        if isinstance(input_data, dict):
            return input_data.copy()
        msg = (
            f"SubWorkflowStep '{self.step_id}' cannot derive child input from "
            f"{type(input_data)!r}; provide input_mapper"
        )
        raise RuntimeError(msg)

    def _extract_child_output(
        self,
        *,
        execution: WorkflowExecution,
        child_workflow: Workflow,
        context: WorkflowContext,
    ) -> Any:
        if self._output_mapper is not None:
            return self._output_mapper(execution, child_workflow, context)

        if len(child_workflow.end_step_ids) != 1:
            msg = (
                f"SubWorkflowStep '{self.step_id}' requires exactly one child end step "
                "for default output extraction; provide output_mapper for custom behavior"
            )
            raise RuntimeError(msg)

        end_step_id = child_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"Child workflow end step '{end_step_id}' did not produce output "
                f"for SubWorkflowStep '{self.step_id}'"
            )
            raise RuntimeError(msg)
        return step_execution.output_data

    def _coerce_output(self, value: Any) -> Any:
        if isinstance(value, self.output_type):
            return value
        candidate = value.model_dump() if isinstance(value, BaseModel) else value
        try:
            return self.output_type.model_validate(candidate)
        except Exception as exc:
            msg = (
                f"SubWorkflowStep '{self.step_id}' could not coerce child output "
                f"of type {type(value)!r} into '{self.output_type.__name__}'. "
                "Provide an output_mapper or return a compatible model/dict."
            )
            raise RuntimeError(msg) from exc

    async def _run_child_workflow(
        self,
        *,
        child_workflow: Workflow,
        initial_input: dict[str, Any] | None = None,
        checkpoint: WorkflowCheckpoint | None = None,
        responses: dict[str, Any] | None = None,
    ) -> tuple[WorkflowExecution, WorkflowCheckpoint | None]:
        from agentbyte.workflow.core.runner import WorkflowRunner

        checkpoint_store = InMemoryCheckpointStore()
        checkpoint_config = CheckpointConfig(
            store=checkpoint_store,
            auto_save=True,
            save_interval_steps=1,
        )
        execution = await WorkflowRunner(
            max_concurrent_steps=self._child_max_concurrent_steps
        ).run(
            child_workflow,
            initial_input,
            checkpoint=checkpoint,
            checkpoint_config=checkpoint_config,
            responses=responses,
        )
        latest_checkpoint = await checkpoint_store.load_latest(child_workflow.id)
        return execution, latest_checkpoint

    async def _propagate_child_request(
        self,
        *,
        context: WorkflowContext,
        invocation_id: str,
        child_checkpoint: WorkflowCheckpoint,
        child_execution: WorkflowExecution,
    ) -> None:
        child_requests = list(child_execution.pending_requests.values())
        if not child_requests:
            msg = "Sub-workflow suspended without a pending request to propagate"
            raise RuntimeError(msg)
        if len(child_requests) > 1:
            msg = (
                f"SubWorkflowStep '{self.step_id}' does not yet support propagating "
                f"multiple concurrent child pending requests (got {len(child_requests)})"
            )
            raise RuntimeError(msg)
        child_request = child_requests[0]

        subworkflow_metadata = {
            **child_request.metadata,
            "subworkflow": {
                "invocation_id": invocation_id,
                "child_workflow_id": child_checkpoint.workflow_id,
                "child_request_id": child_request.request_id,
                "child_request_key": child_request.key,
            },
        }

        try:
            if child_request.kind == WorkflowRequestKind.APPROVAL:
                await context.request_approval(
                    key=child_request.key,
                    reason=child_request.reason or f"Sub-workflow approval for '{child_request.key}'",
                    payload=child_request.payload,
                    response_type=Any,
                    metadata=subworkflow_metadata,
                )
            else:
                await context.request_input(
                    key=child_request.key,
                    response_type=Any,
                    prompt=child_request.prompt or f"Sub-workflow input for '{child_request.key}'",
                    metadata=subworkflow_metadata,
                )
        except _WorkflowRequestSuspended as parent_signal:
            self._request_routes[parent_signal.request.request_id] = {
                "invocation_id": invocation_id,
                "child_request_id": child_request.request_id,
            }
            raise


__all__ = ["SubWorkflowStep"]
