<?php
namespace boru\ai\Stream\Listener;

use boru\ai\Stream\StreamEvent;
use boru\ai\Stream\StreamEventListenerInterface;

class FunctionCallListener implements StreamEventListenerInterface
{
    /** @var array */
    protected $events = array();

    /** @var array */
    protected $calls = array();

    /** @var array */
    protected $startedCallbacks = array();

    /** @var array */
    protected $completedCallbacks = array();

    /** @var array */
    protected $failedCallbacks = array();

    public function onStarted($functionName, $callback)
    {
        return $this->registerCallback($this->startedCallbacks, $functionName, $callback);
    }

    public function onCompleted($functionName, $callback)
    {
        return $this->registerCallback($this->completedCallbacks, $functionName, $callback);
    }

    public function onFailed($functionName, $callback)
    {
        return $this->registerCallback($this->failedCallbacks, $functionName, $callback);
    }

    public function handle(StreamEvent $event)
    {
        $type = $event->type();
        if (!in_array($type, array(
            StreamEvent::RESPONSE_TOOL_CALL_STARTED,
            StreamEvent::RESPONSE_TOOL_CALL_COMPLETED,
            StreamEvent::TOOL_EXECUTION_STARTED,
            StreamEvent::TOOL_EXECUTION_COMPLETED,
            StreamEvent::TOOL_EXECUTION_FAILED,
        ), true)) {
            return;
        }

        $this->events[] = $event;
        $toolCallId = $event->toolCallId();
        if (!$toolCallId) {
            return;
        }

        if (!isset($this->calls[$toolCallId])) {
            $this->calls[$toolCallId] = array(
                'tool_call_id' => $toolCallId,
                'tool_name' => $event->toolName(),
                'status' => null,
                'started' => false,
                'completed' => false,
                'failed' => false,
                'arguments' => null,
                'arguments_json' => null,
                'output' => null,
                'metadata' => array(),
            );
        }

        $call = &$this->calls[$toolCallId];
        if ($event->toolName()) {
            $call['tool_name'] = $event->toolName();
        }

        if ($type === StreamEvent::RESPONSE_TOOL_CALL_STARTED || $type === StreamEvent::TOOL_EXECUTION_STARTED) {
            $call['started'] = true;
            $call['status'] = 'started';
            $call['arguments'] = $event->data('arguments');
            $call['arguments_json'] = $event->data('arguments_json');
            $this->fireCallbacks($this->startedCallbacks, $event, $call);
        } elseif ($type === StreamEvent::RESPONSE_TOOL_CALL_COMPLETED || $type === StreamEvent::TOOL_EXECUTION_COMPLETED) {
            $call['completed'] = true;
            $call['status'] = 'completed';
            $call['output'] = $event->data('output');
            $call['metadata'] = $event->data('metadata');
            $this->fireCallbacks($this->completedCallbacks, $event, $call);
        } elseif ($type === StreamEvent::TOOL_EXECUTION_FAILED) {
            $call['failed'] = true;
            $call['status'] = 'failed';
            $call['output'] = $event->data('output');
            $call['metadata'] = $event->data('metadata');
            $this->fireCallbacks($this->failedCallbacks, $event, $call);
        }
        unset($call);
    }

    protected function registerCallback(array &$target, $functionName, $callback)
    {
        if (!is_callable($callback)) {
            throw new \InvalidArgumentException('Function call listener callback must be callable');
        }
        if ($functionName === null || $functionName === '') {
            $functionName = '*';
        }
        if (!isset($target[$functionName])) {
            $target[$functionName] = array();
        }
        $target[$functionName][] = $callback;
        return $this;
    }

    protected function fireCallbacks(array $callbacksByName, StreamEvent $event, array $call)
    {
        $functionName = isset($call['tool_name']) ? $call['tool_name'] : 'unknown';

        if (isset($callbacksByName[$functionName])) {
            foreach ($callbacksByName[$functionName] as $callback) {
                call_user_func($callback, $event, $functionName, $call['arguments'], $call['output']);
            }
        }

        if (isset($callbacksByName['*'])) {
            foreach ($callbacksByName['*'] as $callback) {
                call_user_func($callback, $event, $functionName, $call['arguments'], $call['output']);
            }
        }
    }

    public function events()
    {
        return $this->events;
    }

    public function calls()
    {
        return $this->calls;
    }
}
