<?php

declare(strict_types=1);

namespace WordPress\XaiAiProvider\Metadata;

use WordPress\AiClient\Messages\Enums\ModalityEnum;
use WordPress\AiClient\Providers\Http\DTO\Request;
use WordPress\AiClient\Providers\Http\DTO\Response;
use WordPress\AiClient\Providers\Http\Enums\HttpMethodEnum;
use WordPress\AiClient\Providers\Http\Exception\ResponseException;
use WordPress\AiClient\Providers\Models\DTO\ModelMetadata;
use WordPress\AiClient\Providers\Models\DTO\SupportedOption;
use WordPress\AiClient\Providers\Models\Enums\CapabilityEnum;
use WordPress\AiClient\Providers\Models\Enums\OptionEnum;
use WordPress\AiClient\Providers\OpenAiCompatibleImplementation\AbstractOpenAiCompatibleModelMetadataDirectory;
use WordPress\XaiAiProvider\Provider\XaiProvider;

/**
 * Class for the xAI (Grok) model metadata directory.
 *
 * @since 1.0.0
 *
 * @phpstan-type ModelsResponseData array{
 *     data: list<array{id: string}>
 * }
 */
class XaiModelMetadataDirectory extends AbstractOpenAiCompatibleModelMetadataDirectory
{
    /**
     * {@inheritDoc}
     *
     * @since 1.0.0
     */
    protected function createRequest(HttpMethodEnum $method, string $path, array $headers = [], $data = null): Request
    {
        return new Request(
            $method,
            XaiProvider::url($path),
            $headers,
            $data
        );
    }

    /**
     * {@inheritDoc}
     *
     * @since 1.0.0
     */
    protected function parseResponseToModelMetadataList(Response $response): array
    {
        /** @var ModelsResponseData $responseData */
        $responseData = $response->getData();
        if (!isset($responseData['data']) || !$responseData['data']) {
            throw ResponseException::fromMissingData('xAI', 'data');
        }

        $grokCapabilities = [
            CapabilityEnum::textGeneration(),
            CapabilityEnum::chatHistory(),
        ];

        $grokBaseOptions = [
            new SupportedOption(OptionEnum::systemInstruction()),
            new SupportedOption(OptionEnum::candidateCount()),
            new SupportedOption(OptionEnum::maxTokens()),
            new SupportedOption(OptionEnum::temperature()),
            new SupportedOption(OptionEnum::topP()),
            new SupportedOption(OptionEnum::stopSequences()),
            new SupportedOption(OptionEnum::presencePenalty()),
            new SupportedOption(OptionEnum::frequencyPenalty()),
            new SupportedOption(OptionEnum::logprobs()),
            new SupportedOption(OptionEnum::topLogprobs()),
            new SupportedOption(OptionEnum::outputMimeType(), ['text/plain', 'application/json']),
            new SupportedOption(OptionEnum::outputSchema()),
            new SupportedOption(OptionEnum::functionDeclarations()),
            new SupportedOption(OptionEnum::customOptions()),
        ];

        $grokTextOptions = array_merge($grokBaseOptions, [
            new SupportedOption(OptionEnum::inputModalities(), [[ModalityEnum::text()]]),
            new SupportedOption(OptionEnum::outputModalities(), [[ModalityEnum::text()]]),
        ]);

        $grokMultimodalInputOptions = array_merge($grokBaseOptions, [
            new SupportedOption(
                OptionEnum::inputModalities(),
                [
                    [ModalityEnum::text()],
                    [ModalityEnum::text(), ModalityEnum::image()],
                ]
            ),
            new SupportedOption(OptionEnum::outputModalities(), [[ModalityEnum::text()]]),
        ]);

        $modelsData = (array) $responseData['data'];

        $models = array_values(
            array_map(
                static function (array $modelData) use (
                    $grokCapabilities,
                    $grokTextOptions,
                    $grokMultimodalInputOptions
                ): ModelMetadata {
                    $modelId = $modelData['id'];

                    if (self::supportsMultimodalTextInput($modelId)) {
                        $modelOptions = $grokMultimodalInputOptions;
                    } else {
                        $modelOptions = $grokTextOptions;
                    }

                    return new ModelMetadata(
                        $modelId,
                        $modelId, // The xAI API does not return a display name.
                        $grokCapabilities,
                        $modelOptions
                    );
                },
                $modelsData
            )
        );

        usort($models, [$this, 'modelSortCallback']);

        return $models;
    }

    /**
     * Checks whether an xAI text generation model supports multimodal (image) input.
     *
     * @since 1.0.0
     *
     * @param string $modelId The model ID.
     * @return bool True if the model supports multimodal text input, false otherwise.
     */
    private static function supportsMultimodalTextInput(string $modelId): bool
    {
        return str_contains($modelId, 'vision') || (bool) preg_match('/^grok-4/', $modelId);
    }

    /**
     * Callback function for sorting models by ID, to be used with `usort()`.
     *
     * This method expresses preferences for certain models or model families within the provider by putting them
     * earlier in the sorted list. The objective is not to be opinionated about which models are better, but to ensure
     * that more commonly used, more recent, or flagship models are presented first to users.
     *
     * @since 1.0.0
     *
     * @param ModelMetadata $a First model.
     * @param ModelMetadata $b Second model.
     * @return int Comparison result.
     */
    protected function modelSortCallback(ModelMetadata $a, ModelMetadata $b): int
    {
        $aId = $a->getId();
        $bId = $b->getId();

        // Prefer non-preview, non-beta models over preview/beta ones.
        foreach (['-preview', '-beta'] as $suffix) {
            if (str_contains($aId, $suffix) && !str_contains($bId, $suffix)) {
                return 1;
            }
            if (str_contains($bId, $suffix) && !str_contains($aId, $suffix)) {
                return -1;
            }
        }

        // Prefer 'grok-' models over any other model families.
        if (str_starts_with($aId, 'grok-') && !str_starts_with($bId, 'grok-')) {
            return -1;
        }
        if (str_starts_with($bId, 'grok-') && !str_starts_with($aId, 'grok-')) {
            return 1;
        }

        // Prefer 'grok' models with version numbers (e.g. 'grok-4', 'grok-3') over those without.
        $aMatch = preg_match('/^grok-([0-9.]+)(-[a-z0-9-]+)?$/', $aId, $aMatches);
        $bMatch = preg_match('/^grok-([0-9.]+)(-[a-z0-9-]+)?$/', $bId, $bMatches);
        if ($aMatch && !$bMatch) {
            return -1;
        }
        if ($bMatch && !$aMatch) {
            return 1;
        }
        if ($aMatch && $bMatch) {
            // Prefer later model versions.
            $aVersion = $aMatches[1];
            $bVersion = $bMatches[1];
            if (version_compare($aVersion, $bVersion, '>')) {
                return -1;
            }
            if (version_compare($bVersion, $aVersion, '>')) {
                return 1;
            }

            // Prefer models without a suffix (i.e. base models) over those with a suffix.
            if (!isset($aMatches[2]) && isset($bMatches[2])) {
                return -1;
            }
            if (!isset($bMatches[2]) && isset($aMatches[2])) {
                return 1;
            }
        }

        // Fallback: Sort alphabetically.
        return strcmp($aId, $bId);
    }
}