<?php

declare(strict_types=1);

namespace WordPress\NvidiaAiProvider\Metadata;

use WordPress\AiClient\Files\Enums\FileTypeEnum;
use WordPress\AiClient\Files\Enums\MediaOrientationEnum;
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\NvidiaAiProvider\Provider\NvidiaProvider;

/**
 * Class for the NVIDIA NIM model metadata directory.
 *
 * The NVIDIA NIM `/models` endpoint returns the full public model catalog (every model NVIDIA has
 * ever published), regardless of the API key. Only a subset of those models is actually provisioned
 * for inference on a given account; the rest return `404 "Function not found for account"` when
 * called. The endpoint also reports no capability information.
 *
 * To guarantee that every model surfaced to WordPress is actually callable, this directory exposes a
 * curated allowlist of models that have been verified against the NVIDIA NIM `chat/completions`
 * endpoint. The live catalog response is still requested (which doubles as the provider availability
 * / API-key check), and the allowlist is intersected with it, so a model that NVIDIA removes from the
 * catalog automatically disappears here too.
 *
 * Each allowlisted model is classified as a plain text (chat) model, a vision-capable chat model that
 * additionally accepts image input, or an image generation model.
 *
 * Image generation (the FLUX family) is served by NVIDIA's separate GenAI endpoint
 * (`ai.api.nvidia.com/v1/genai`) rather than the OpenAI-compatible `integrate.api.nvidia.com`
 * endpoint, and those models do not appear in the `/v1/models` catalog. They are therefore surfaced
 * unconditionally rather than intersected with the catalog.
 *
 * @since 1.0.0
 *
 * @phpstan-type ModelsResponseData array{
 *     data: list<array{id: string}>
 * }
 */
class NvidiaModelMetadataDirectory extends AbstractOpenAiCompatibleModelMetadataDirectory
{
    /**
     * Curated allowlist of NVIDIA NIM models that are verified callable for generation.
     *
     * Keys are model IDs, values are the model kind: `'text'` for chat models, `'vision'` for chat
     * models that additionally accept image input, or `'image'` for image generation models. The order
     * here is the order in which models are presented to users, with broadly useful, reliable models
     * first (the first entry is what the WordPress AI plugin auto-selects when no specific model is
     * requested).
     *
     * Models that NVIDIA exposes in the catalog but does not reliably provision for inference (they
     * return 404, or cold-start far beyond the WordPress HTTP timeout) are intentionally excluded.
     *
     * @since 1.0.0
     *
     * @var array<string, string>
     */
    private const AVAILABLE_MODELS = [
        // Meta Llama — fast, reliable, broadly available general-purpose models.
        'meta/llama-3.3-70b-instruct' => 'text',
        'meta/llama-3.1-8b-instruct' => 'text',
        'meta/llama-3.1-70b-instruct' => 'text',
        'meta/llama-3.2-3b-instruct' => 'text',
        'meta/llama-3.2-1b-instruct' => 'text',
        'meta/llama-4-maverick-17b-128e-instruct' => 'text',
        // NVIDIA Nemotron — NVIDIA's own model family.
        'nvidia/nemotron-3-super-120b-a12b' => 'text',
        'nvidia/nemotron-3-nano-30b-a3b' => 'text',
        'nvidia/nemotron-3-nano-omni-30b-a3b-reasoning' => 'text',
        'nvidia/nemotron-3-ultra-550b-a55b' => 'text',
        'nvidia/nvidia-nemotron-nano-9b-v2' => 'text',
        'nvidia/nemotron-mini-4b-instruct' => 'text',
        'nvidia/llama-3.3-nemotron-super-49b-v1.5' => 'text',
        'nvidia/llama-3.3-nemotron-super-49b-v1' => 'text',
        'nvidia/ising-calibration-1-35b-a3b' => 'text',
        'nvidia/riva-translate-4b-instruct-v1.1' => 'text',
        // Mistral.
        'mistralai/mistral-large-3-675b-instruct-2512' => 'text',
        'mistralai/mistral-medium-3.5-128b' => 'text',
        'mistralai/mistral-small-4-119b-2603' => 'text',
        'mistralai/mistral-nemotron' => 'text',
        'mistralai/ministral-14b-instruct-2512' => 'text',
        'mistralai/mixtral-8x7b-instruct-v0.1' => 'text',
        // OpenAI open-weight models hosted on NVIDIA NIM.
        'openai/gpt-oss-120b' => 'text',
        'openai/gpt-oss-20b' => 'text',
        // Qwen.
        'qwen/qwen3.5-122b-a10b' => 'text',
        'qwen/qwen3-next-80b-a3b-instruct' => 'text',
        // Microsoft Phi.
        'microsoft/phi-4-mini-instruct' => 'text',
        // Google Gemma.
        'google/gemma-3n-e2b-it' => 'text',
        'google/gemma-2-2b-it' => 'text',
        // Other verified general-purpose models.
        'moonshotai/kimi-k2.6' => 'text',
        'minimaxai/minimax-m3' => 'text',
        'minimaxai/minimax-m2.7' => 'text',
        'stepfun-ai/step-3.7-flash' => 'text',
        'stepfun-ai/step-3.5-flash' => 'text',
        'z-ai/glm-5.1' => 'text',
        'abacusai/dracarys-llama-3.1-70b-instruct' => 'text',
        'bytedance/seed-oss-36b-instruct' => 'text',
        'stockmark/stockmark-2-100b-instruct' => 'text',
        'sarvamai/sarvam-m' => 'text',
        'upstage/solar-10.7b-instruct' => 'text',
        // Vision-capable chat models (accept image input in addition to text).
        'meta/llama-3.2-90b-vision-instruct' => 'vision',
        'meta/llama-3.2-11b-vision-instruct' => 'vision',
        'nvidia/nemotron-nano-12b-v2-vl' => 'vision',
        'nvidia/llama-3.1-nemotron-nano-vl-8b-v1' => 'vision',
        'microsoft/phi-4-multimodal-instruct' => 'vision',
        // Image generation (FLUX) — served by the NVIDIA GenAI endpoint, not in the /v1/models catalog.
        'black-forest-labs/flux.1-dev' => 'image',
        'black-forest-labs/flux.1-schnell' => 'image',
        'black-forest-labs/flux.2-klein-4b' => 'image',
    ];

    /**
     * {@inheritDoc}
     *
     * @since 1.0.0
     */
    protected function createRequest(HttpMethodEnum $method, string $path, array $headers = [], $data = null): Request
    {
        return new Request(
            $method,
            NvidiaProvider::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('NVIDIA', 'data');
        }

        /*
         * The NVIDIA NIM API does not return model capabilities, so they are hardcoded here.
         * Chat models support text generation and chat history; vision-capable models additionally
         * accept image input.
         */
        $chatCapabilities = [
            CapabilityEnum::textGeneration(),
            CapabilityEnum::chatHistory(),
        ];
        $textOptions = [
            new SupportedOption(OptionEnum::systemInstruction()),
            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::outputMimeType(), ['text/plain', 'application/json']),
            new SupportedOption(OptionEnum::outputSchema()),
            new SupportedOption(OptionEnum::functionDeclarations()),
            new SupportedOption(OptionEnum::customOptions()),
            new SupportedOption(OptionEnum::inputModalities(), [[ModalityEnum::text()]]),
            new SupportedOption(OptionEnum::outputModalities(), [[ModalityEnum::text()]]),
        ];
        $visionOptions = array_merge(
            array_slice($textOptions, 0, -2),
            [
                new SupportedOption(
                    OptionEnum::inputModalities(),
                    [
                        [ModalityEnum::text()],
                        [ModalityEnum::text(), ModalityEnum::image()],
                    ]
                ),
                new SupportedOption(OptionEnum::outputModalities(), [[ModalityEnum::text()]]),
            ]
        );

        // Image generation (FLUX) capabilities and options.
        $imageCapabilities = [
            CapabilityEnum::imageGeneration(),
        ];
        $imageOptions = [
            new SupportedOption(OptionEnum::candidateCount()),
            // The NVIDIA GenAI image endpoint returns inline JPEG image data.
            new SupportedOption(OptionEnum::outputMimeType(), ['image/jpeg']),
            new SupportedOption(OptionEnum::outputFileType(), [FileTypeEnum::inline()]),
            new SupportedOption(OptionEnum::outputMediaOrientation(), [
                MediaOrientationEnum::square(),
                MediaOrientationEnum::landscape(),
                MediaOrientationEnum::portrait(),
            ]),
            new SupportedOption(OptionEnum::outputMediaAspectRatio(), ['1:1', '3:2', '2:3']),
            new SupportedOption(OptionEnum::customOptions()),
            new SupportedOption(OptionEnum::inputModalities(), [[ModalityEnum::text()]]),
            new SupportedOption(OptionEnum::outputModalities(), [[ModalityEnum::image()]]),
        ];

        // Collect the set of model IDs that NVIDIA currently reports in its catalog.
        $catalogIds = [];
        foreach ((array) $responseData['data'] as $modelData) {
            if (isset($modelData['id']) && is_string($modelData['id'])) {
                $catalogIds[$modelData['id']] = true;
            }
        }

        /*
         * Surface allowlisted models in the curated order. Text and vision (chat) models are only
         * included if they are still present in the live catalog, which guarantees they are verified
         * callable and stays in sync with NVIDIA removing models. Image generation models live on a
         * separate endpoint and are not in this catalog, so they are surfaced unconditionally.
         */
        $models = [];
        foreach (self::AVAILABLE_MODELS as $modelId => $kind) {
            if ('image' === $kind) {
                $models[] = new ModelMetadata($modelId, $modelId, $imageCapabilities, $imageOptions);
                continue;
            }

            if (!isset($catalogIds[$modelId])) {
                continue;
            }

            $models[] = new ModelMetadata(
                $modelId,
                $modelId, // The NVIDIA NIM API does not return a display name.
                $chatCapabilities,
                'vision' === $kind ? $visionOptions : $textOptions
            );
        }

        return $models;
    }
}