<?php
namespace boru\ai\Provider\Gemini;

use boru\ai\Provider\ProviderSignal;
use boru\ai\Schema\Response;

class GeminiStreamParser
{
    /** @var GeminiMapper */
    protected $mapper;

    /** @var string */
    protected $buffer = '';

    /** @var string */
    protected $currentText = '';

    /** @var Response|null */
    protected $finalResponse;

    /** @var bool */
    protected $completed = false;

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

    /** @var array|null */
    protected $error;

    public function __construct(GeminiMapper $mapper = null)
    {
        $this->mapper = $mapper ? $mapper : new GeminiMapper();
    }

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

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

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

    public function push($chunk, $signalCallback = null)
    {
        $normalized = str_replace("\r\n", "\n", (string)$chunk);
        $this->buffer .= $normalized;

        while (($pos = strpos($this->buffer, "\n\n")) !== false) {
            $rawEvent = substr($this->buffer, 0, $pos);
            $this->buffer = substr($this->buffer, $pos + 2);
            $this->parseEventBlock($rawEvent, $signalCallback);
        }
    }

    protected function parseEventBlock($rawEvent, $signalCallback = null)
    {
        $lines = preg_split('/\r?\n/', trim($rawEvent));
        $dataLines = array();

        foreach ($lines as $line) {
            if (strpos($line, 'data:') === 0) {
                $dataLines[] = trim(substr($line, 5));
            }
        }

        if (empty($dataLines)) {
            return;
        }

        $data = implode("\n", $dataLines);
        if ($data === '[DONE]') {
            $this->emitTextCompleted($signalCallback);
            return;
        }

        $decoded = json_decode($data, true);
        if (!is_array($decoded)) {
            return;
        }

        if (isset($decoded['error'])) {
            $this->error = $decoded['error'];
            $signal = new ProviderSignal(ProviderSignal::RESPONSE_ERROR, array(
                'error' => $decoded['error'],
            ), $this->finalResponse);
            $this->dispatch($signal, $signalCallback);
            return;
        }

        $parts = $this->extractParts($decoded);
        foreach ($parts as $part) {
            if (isset($part['text'])) {
                $delta = (string)$part['text'];
                if ($delta !== '') {
                    $this->currentText .= $delta;
                    $signal = new ProviderSignal(ProviderSignal::TEXT_DELTA, array(
                        'delta' => $delta,
                        'text' => $delta,
                    ));
                    $this->dispatch($signal, $signalCallback);
                }
            }

            if (isset($part['functionCall']) && is_array($part['functionCall'])) {
                $this->emitFunctionCallSignals($part['functionCall'], $signalCallback, $decoded);
            }
        }

        $this->finalResponse = $this->mapper->mapResponse($decoded);
        if ($this->finalResponse instanceof Response && $this->currentText !== '') {
            $this->finalResponse->text($this->currentText);
        }

        $finishReason = isset($decoded['candidates'][0]['finishReason']) ? $decoded['candidates'][0]['finishReason'] : null;
        $finishMessage = isset($decoded['candidates'][0]['finishMessage']) ? $decoded['candidates'][0]['finishMessage'] : null;
        if ($finishReason !== null) {
            $this->completed = true;
            $signal = new ProviderSignal(ProviderSignal::RESPONSE_COMPLETED, array(
                'finish_reason' => $finishReason,
                'finish_message' => $finishMessage,
            ), $this->finalResponse);
            $this->dispatch($signal, $signalCallback);
            $this->emitTextCompleted($signalCallback);
        }
    }

    protected function emitFunctionCallSignals(array $functionCall, $signalCallback = null, array $payload = array())
    {
        $toolCallId = isset($functionCall['id']) ? $functionCall['id'] : (isset($functionCall['name']) ? $functionCall['name'] : null);
        $toolName = isset($functionCall['name']) ? $functionCall['name'] : null;
        $arguments = isset($functionCall['args']) && is_array($functionCall['args']) ? $functionCall['args'] : array();
        $argumentsJson = json_encode($arguments);

        if (!$toolCallId || !$toolName) {
            return;
        }

        if (!isset($this->emittedToolCalls[$toolCallId])) {
            $started = new ProviderSignal(ProviderSignal::TOOL_CALL_STARTED, array(
                'tool_call_id' => $toolCallId,
                'tool_name' => $toolName,
                'arguments' => $arguments,
                'arguments_json' => $argumentsJson,
            ), $this->finalResponse);
            $this->dispatch($started, $signalCallback);
            $this->emittedToolCalls[$toolCallId] = true;
        }

        $argumentsSignal = new ProviderSignal(ProviderSignal::TOOL_CALL_ARGUMENTS_DELTA, array(
            'tool_call_id' => $toolCallId,
            'tool_name' => $toolName,
            'arguments' => $arguments,
            'arguments_json' => $argumentsJson,
        ), $this->finalResponse);
        $this->dispatch($argumentsSignal, $signalCallback);

        $completed = new ProviderSignal(ProviderSignal::TOOL_CALL_COMPLETED, array(
            'tool_call_id' => $toolCallId,
            'tool_name' => $toolName,
            'arguments' => $arguments,
            'arguments_json' => $argumentsJson,
            'finish_reason' => isset($payload['candidates'][0]['finishReason']) ? $payload['candidates'][0]['finishReason'] : null,
            'finish_message' => isset($payload['candidates'][0]['finishMessage']) ? $payload['candidates'][0]['finishMessage'] : null,
        ), $this->finalResponse);
        $this->dispatch($completed, $signalCallback);
    }

    protected function emitTextCompleted($signalCallback = null)
    {
        if ($this->currentText === '') {
            return;
        }

        if ($this->finalResponse instanceof Response) {
            $this->finalResponse->text($this->currentText);
        }

        $signal = new ProviderSignal(ProviderSignal::TEXT_COMPLETED, array(
            'text' => $this->currentText,
        ), $this->finalResponse);
        $this->dispatch($signal, $signalCallback);
        $this->currentText = '';
    }

    protected function extractParts(array $payload)
    {
        if (!isset($payload['candidates']) || !is_array($payload['candidates'])) {
            return array();
        }

        $parts = array();
        foreach ($payload['candidates'] as $candidate) {
            if (!is_array($candidate) || !isset($candidate['content']['parts']) || !is_array($candidate['content']['parts'])) {
                continue;
            }

            foreach ($candidate['content']['parts'] as $part) {
                $parts[] = $part;
            }
        }

        return $parts;
    }

    protected function dispatch(ProviderSignal $signal, $signalCallback = null)
    {
        if (is_callable($signalCallback)) {
            call_user_func($signalCallback, $signal);
        }
    }
}
