<?php

declare(strict_types=1);

namespace NeuronAI\Agent\Middleware;

use NeuronAI\Agent\Events\ToolCallEvent;
use NeuronAI\Agent\Nodes\ToolNode;
use NeuronAI\Agent\Tools\ToolRejectionHandler;
use NeuronAI\Tools\ToolInterface;
use NeuronAI\Workflow\Events\Event;
use NeuronAI\Workflow\Interrupt\Action;
use NeuronAI\Workflow\Interrupt\ActionDecision;
use NeuronAI\Workflow\Interrupt\ApprovalRequest;
use NeuronAI\Workflow\Interrupt\WorkflowInterrupt;
use NeuronAI\Workflow\Middleware\WorkflowMiddleware;
use NeuronAI\Workflow\NodeInterface;
use NeuronAI\Workflow\WorkflowState;

use function array_filter;
use function count;
use function is_callable;
use function is_int;
use function is_string;
use function json_encode;
use function sprintf;
use function uniqid;

use const JSON_PRETTY_PRINT;

class ToolApproval implements WorkflowMiddleware
{
    /**
     * @param array<int|string, string|callable(array): bool> $tools Tools that require approval.
     *   - Empty array: all tools require approval (default)
     *   - Numeric key + string value: tool name or class-string always requires approval
     *   - String key + callable value: tool name or class-string with conditional callback.
     *     The callback receives the tool's input arguments (array) and returns bool
     *     (true = requires approval, false = skip).
     *
     * Example:
     *   new ToolApproval([
     *       'delete_file',
     *       'transfer_money' => fn(array $args) => $args['amount'] > 100,
     *   ])
     */
    public function __construct(
        protected array $tools = []
    ) {
    }

    /**
     * Execute before the node runs.
     *
     * On initial run: Inspects tools and creates interrupt request for approval.
     * On resume: Processes human decisions and modifies tools accordingly.
     *
     * @param ToolNode $node
     * @param ToolCallEvent $event
     * @throws WorkflowInterrupt
     */
    public function before(NodeInterface $node, Event $event, WorkflowState $state): void
    {
        if (!$event instanceof ToolCallEvent) {
            return;
        }

        // Check if we're resuming
        if ($node->isResuming() && $node->getResumeRequest() instanceof ApprovalRequest) {
            $this->processDecisions($node->getResumeRequest(), $event);
            return;
        }

        // Initial run: Check if any tools require approval
        $toolsToApprove = $this->filterToolsRequiringApproval($event->toolCallMessage->getTools());

        if ($toolsToApprove === []) {
            // No tools require approval, continue execution
            return;
        }

        // Create the interrupt request with actions for each tool
        $actions = [];
        foreach ($toolsToApprove as $tool) {
            $actions[] = $this->createAction($tool);
        }

        $count = count($actions);
        $interruptRequest = new ApprovalRequest(
            message: sprintf(
                '%d tool call%s require%s approval before execution',
                $count,
                $count === 1 ? '' : 's',
                $count === 1 ? 's' : ''
            ),
            actions: $actions
        );

        throw new WorkflowInterrupt($interruptRequest, $node, $state, $event);
    }

    /**
     * Execute after the node runs.
     */
    public function after(NodeInterface $node, Event $result, WorkflowState $state): void
    {
        //
    }

    /**
     * Filter tools that require approval based on configuration.
     *
     * @param ToolInterface[] $tools
     * @return ToolInterface[]
     */
    protected function filterToolsRequiringApproval(array $tools): array
    {
        if ($this->tools === []) {
            // Empty array means all tools require approval
            return $tools;
        }

        return array_filter(
            $tools,
            $this->toolRequiresApproval(...)
        );
    }

    /**
     * Determine if a specific tool requires approval.
     *
     * Checks the tool against the configured tools list, handling both
     * unconditional (string) and conditional (callable) entries.
     */
    protected function toolRequiresApproval(ToolInterface $tool): bool
    {
        $toolName = $tool->getName();
        $toolClass = $tool::class;

        foreach ($this->tools as $key => $value) {
            // Numeric key + string value: unconditional approval
            if (is_int($key) && is_string($value) && ($value === $toolName || $value === $toolClass)) {
                return true;
            }

            // String key + callable value: conditional approval
            if (is_string($key) && is_callable($value) && ($key === $toolName || $key === $toolClass)) {
                return $value($tool->getInputs());
            }
        }

        return false;
    }

    /**
     * Create an Action for a tool that requires approval.
     */
    protected function createAction(ToolInterface $tool): Action
    {
        $inputs = $tool->getInputs();
        $inputsDescription = $inputs === []
            ? '(no arguments)'
            : json_encode($inputs, JSON_PRETTY_PRINT);

        return new Action(
            id: $tool->getCallId() ?? uniqid('tool_'),
            name: $tool->getName(),
            description: $inputsDescription,
            decision: ActionDecision::Pending
        );
    }

    /**
     * Process human decisions and modify tools accordingly.
     *
     * This method modifies the tools in-place based on human decisions:
     *  - Rejected: Tool callback is replaced to return rejection message
     *  - Edited: Tool inputs are modified
     *  - Approved: No changes, tool executes normally
     */
    protected function processDecisions(
        ApprovalRequest $request,
        ToolCallEvent   $event,
    ): void {
        foreach ($event->toolCallMessage->getTools() as $tool) {
            $toolCallId = $tool->getCallId();
            if ($toolCallId === null) {
                // Tool doesn't require approval, skip
                continue;
            }

            $action = $request->getAction($toolCallId);

            if (!$action instanceof Action) {
                // Tool doesn't require approval, skip
                continue;
            }

            // Process based on decision
            if ($action->isRejected()) {
                $this->handleRejectedTool($tool, $action);
            }

            // If approved, do nothing - the tool will be executed normally
        }
    }

    /**
     * Handle a rejected tool by replacing its callback with a rejection message.
     *
     * This prevents the tool from executing its actual logic and instead
     * returns a human-readable rejection message that the AI can process.
     *
     * Uses ToolRejectionHandler instead of a closure to ensure serializability
     * when workflows are interrupted multiple times.
     */
    protected function handleRejectedTool(ToolInterface $tool, Action $action): void
    {
        $feedback = $action->feedback ?? 'No specific instruction provided.';
        $rejectionMessage = sprintf(
            "TOOL NOT EXECUTED. The user rejected this action. User instruction: %s. Do not attempt this tool again. Follow the user's instruction.",
            $feedback
        );

        // Replace the tool's callback with a serializable rejection handler
        $tool->setCallable(new ToolRejectionHandler($rejectionMessage));
    }
}
