<?php

declare(strict_types=1);

namespace ContextPaging\Middleware;

use Psr\Http\Message\ResponseInterface;
use Psr\Http\Message\ServerRequestInterface;
use Psr\Http\Server\MiddlewareInterface;
use Psr\Http\Server\RequestHandlerInterface;

/**
 * Model Quirks Middleware
 *
 * Handles model-specific behavior adjustments so the core logic stays model-agnostic.
 * Mutates the request payload based on the model being used.
 *
 * Example quirks:
 * - SmolLM3: Inject `/no_think` to suppress reasoning tokens
 * - Some models: Strip tool definitions to avoid confusion
 */
class ModelQuirksMiddleware implements MiddlewareInterface
{
    /**
     * Model quirk configurations.
     * Keys are regex patterns to match model names.
     */
    private const QUIRK_CONFIGS = [
        // SmolLM3 - small model that outputs reasoning tokens and gets confused by tools
        '/SmolLM3/i' => [
            'no_think' => true,           // Inject /no_think into system prompt
            'strip_tools' => true,        // Remove tool definitions (doesn't use them correctly)
            'tool_mode' => 'raw',         // Use raw mode if tools are kept
        ],
        // Add more model quirks here as needed
        // '/Llama-3/i' => [...],
        // '/Qwen/i' => [...],
    ];

    /**
     * Process the request and apply model quirks.
     */
    public function process(ServerRequestInterface $request, RequestHandlerInterface $handler): ResponseInterface
    {
        $body = $request->getParsedBody();

        if (!is_array($body)) {
            return $handler->handle($request);
        }

        $model = $body['model'] ?? '';
        $quirks = $this->getQuirksForModel($model);

        if (empty($quirks)) {
            return $handler->handle($request);
        }

        // Apply quirks to the request body
        $body = $this->applyQuirks($body, $quirks);

        // Store original body for reference
        $request = $request->withAttribute('original_body', $request->getParsedBody());
        $request = $request->withParsedBody($body);

        // Also update the raw body for downstream consumers
        $request = $request->withBody(new \Slim\Psr7\Stream(fopen('php://temp', 'r+')));
        $request->getBody()->write(json_encode($body));
        $request->getBody()->rewind();

        return $handler->handle($request);
    }

    /**
     * Get quirks for a specific model.
     */
    public function getQuirksForModel(string $model): array
    {
        foreach (self::QUIRK_CONFIGS as $pattern => $quirks) {
            if (preg_match($pattern, $model)) {
                return $quirks;
            }
        }

        return [];
    }

    /**
     * Apply quirks to the request body.
     */
    private function applyQuirks(array $body, array $quirks): array
    {
        // Inject /no_think into system prompt
        if ($quirks['no_think'] ?? false) {
            $body = $this->injectNoThink($body);
        }

        // Store quirks for later middleware/handlers
        $body['_quirks'] = $quirks;

        return $body;
    }

    /**
     * Inject /no_think into the system prompt.
     */
    private function injectNoThink(array $body): array
    {
        $messages = $body['messages'] ?? [];

        foreach ($messages as $i => $message) {
            if (($message['role'] ?? null) === 'system') {
                $content = $message['content'] ?? '';

                // Only inject if not already present
                if (!str_contains($content, '/no_think')) {
                    $messages[$i]['content'] = rtrim($content) . ' /no_think';
                }

                $body['messages'] = $messages;
                return $body;
            }
        }

        // No system message found, prepend one
        array_unshift($messages, [
            'role' => 'system',
            'content' => '/no_think',
        ]);

        $body['messages'] = $messages;
        return $body;
    }
}
