<?php

declare(strict_types=1);

namespace WordPress\AiClient\Providers\Models\DTO;

use WordPress\AiClient\Common\AbstractDataTransferObject;
use WordPress\AiClient\Common\Exception\InvalidArgumentException;
use WordPress\AiClient\Messages\DTO\Message;
use WordPress\AiClient\Messages\DTO\MessagePart;
use WordPress\AiClient\Messages\Enums\ModalityEnum;
use WordPress\AiClient\Providers\Models\Enums\CapabilityEnum;
use WordPress\AiClient\Providers\Models\Enums\OptionEnum;

/**
 * Represents requirements that implementing code has for AI model selection.
 *
 * This class defines the capabilities and options that a model must support
 * in order to be considered suitable for the implementing code's needs.
 *
 * @since 0.1.0
 *
 * @phpstan-import-type RequiredOptionArrayShape from RequiredOption
 *
 * @phpstan-type ModelRequirementsArrayShape array{
 *     requiredCapabilities: list<string>,
 *     requiredOptions: list<RequiredOptionArrayShape>
 * }
 *
 * @phpstan-type UnmetModelRequirementsShape array{
 *     capabilities: list<CapabilityEnum>,
 *     options: list<RequiredOption>
 * }
 *
 * @extends AbstractDataTransferObject<ModelRequirementsArrayShape>
 */
class ModelRequirements extends AbstractDataTransferObject
{
    public const KEY_REQUIRED_CAPABILITIES = 'requiredCapabilities';
    public const KEY_REQUIRED_OPTIONS = 'requiredOptions';

    /**
     * @var list<CapabilityEnum> The capabilities that the model must support.
     */
    protected array $requiredCapabilities;

    /**
     * @var list<RequiredOption> The options that the model must support with specific values.
     */
    protected array $requiredOptions;

    /**
     * Constructor.
     *
     * @since 0.1.0
     *
     * @param list<CapabilityEnum> $requiredCapabilities The capabilities that the model must support.
     * @param list<RequiredOption> $requiredOptions The options that the model must support with specific values.
     *
     * @throws InvalidArgumentException If arrays are not lists.
     */
    public function __construct(array $requiredCapabilities, array $requiredOptions)
    {
        if (!array_is_list($requiredCapabilities)) {
            throw new InvalidArgumentException('Required capabilities must be a list array.');
        }

        if (!array_is_list($requiredOptions)) {
            throw new InvalidArgumentException('Required options must be a list array.');
        }

        $this->requiredCapabilities = $requiredCapabilities;
        $this->requiredOptions = $requiredOptions;
    }

    /**
     * Gets the capabilities that the model must support.
     *
     * @since 0.1.0
     *
     * @return list<CapabilityEnum> The required capabilities.
     */
    public function getRequiredCapabilities(): array
    {
        return $this->requiredCapabilities;
    }

    /**
     * Gets the options that the model must support with specific values.
     *
     * @since 0.1.0
     *
     * @return list<RequiredOption> The required options.
     */
    public function getRequiredOptions(): array
    {
        return $this->requiredOptions;
    }

    /**
     * Checks whether the given model metadata meets these requirements.
     *
     * @since 0.2.0
     *
     * @param ModelMetadata $metadata The model metadata to check against.
     * @return bool True if the model meets all requirements, false otherwise.
     */
    public function areMetBy(ModelMetadata $metadata): bool
    {
        $unmetRequirements = $this->getUnmetRequirements($metadata);

        return $unmetRequirements['capabilities'] === [] && $unmetRequirements['options'] === [];
    }

    /**
     * Determines which of these requirements the given model metadata does not meet.
     *
     * Unlike {@see self::areMetBy()}, this method reports the specific capabilities and options that
     * are unsupported, so that calling code can explain why a model is unsuitable.
     *
     * @since 1.4.0
     *
     * @param ModelMetadata $metadata The model metadata to check against.
     * @return UnmetModelRequirementsShape The unsupported capabilities and options. Both lists are empty if the
     *                                     model meets all requirements.
     */
    public function getUnmetRequirements(ModelMetadata $metadata): array
    {
        // Create lookup maps for better performance (instead of nested foreach loops)
        $capabilitiesMap = [];
        foreach ($metadata->getSupportedCapabilities() as $capability) {
            $capabilitiesMap[$capability->value] = $capability;
        }

        $optionsMap = [];
        foreach ($metadata->getSupportedOptions() as $option) {
            $optionsMap[$option->getName()->value] = $option;
        }

        // Collect required capabilities that are not supported, using map lookup
        $unmetCapabilities = [];
        foreach ($this->requiredCapabilities as $requiredCapability) {
            if (!isset($capabilitiesMap[$requiredCapability->value])) {
                $unmetCapabilities[] = $requiredCapability;
            }
        }

        // Collect required options that are either unsupported or unsupported with the required value
        $unmetOptions = [];
        foreach ($this->requiredOptions as $requiredOption) {
            // Use map lookup instead of linear search
            $supportedOption = $optionsMap[$requiredOption->getName()->value] ?? null;

            if ($supportedOption === null || !$supportedOption->isSupportedValue($requiredOption->getValue())) {
                $unmetOptions[] = $requiredOption;
            }
        }

        return [
            'capabilities' => $unmetCapabilities,
            'options' => $unmetOptions,
        ];
    }

    /**
     * Creates ModelRequirements from prompt data and model configuration.
     *
     * @since 0.2.0
     *
     * @param CapabilityEnum $capability The capability the model must support.
     * @param list<Message> $messages The messages in the conversation.
     * @param ModelConfig $modelConfig The model configuration.
     * @return self The created requirements.
     */
    public static function fromPromptData(CapabilityEnum $capability, array $messages, ModelConfig $modelConfig): self
    {
        // Start with base capability
        $capabilities = [$capability];
        $inputModalities = [];

        // Check if we have chat history (multiple messages)
        if (count($messages) > 1) {
            $capabilities[] = CapabilityEnum::chatHistory();
        }

        // Analyze all messages to determine required input modalities
        $hasFunctionMessageParts = false;
        foreach ($messages as $message) {
            foreach ($message->getParts() as $part) {
                $modality = self::inputModalityForPart($part);
                if ($modality !== null) {
                    $inputModalities[] = $modality;
                }

                // Check for function calls/responses (these might require special capabilities)
                if ($part->getType()->isFunctionCall() || $part->getType()->isFunctionResponse()) {
                    $hasFunctionMessageParts = true;
                }
            }
        }

        // Convert ModelConfig to RequiredOptions
        $requiredOptions = self::toRequiredOptions($modelConfig);

        // Add additional options based on message analysis
        if ($hasFunctionMessageParts) {
            $requiredOptions = self::includeInRequiredOptions(
                $requiredOptions,
                new RequiredOption(OptionEnum::functionDeclarations(), true)
            );
        }

        // Add input modalities if we have any inputs
        if (!empty($inputModalities)) {
            // Remove duplicates
            $inputModalities = array_unique($inputModalities, SORT_REGULAR);
            $requiredOptions = self::includeInRequiredOptions(
                $requiredOptions,
                new RequiredOption(OptionEnum::inputModalities(), array_values($inputModalities))
            );
        }

        // Step 6: Return new ModelRequirements
        return new self($capabilities, $requiredOptions);
    }

    /**
     * Creates ModelRequirements from embedding input data and model configuration.
     *
     * Unlike {@see self::fromPromptData()}, embedding inputs are independent items rather than a
     * conversation, so no chat history capability is inferred. Each input contributes its input
     * modality (text or file) to the requirements.
     *
     * @since 1.4.0
     *
     * @param list<MessagePart> $inputs The embedding inputs.
     * @param ModelConfig $modelConfig The model configuration.
     * @return self The created requirements.
     */
    public static function fromEmbeddingData(array $inputs, ModelConfig $modelConfig): self
    {
        $capabilities = [CapabilityEnum::embeddingGeneration()];
        $inputModalities = [];

        // Analyze each input to determine required input modalities. Function call/response parts
        // are not valid embedding inputs and are rejected before reaching this point.
        foreach ($inputs as $part) {
            $modality = self::inputModalityForPart($part);
            if ($modality !== null) {
                $inputModalities[] = $modality;
            }
        }

        // Convert ModelConfig to RequiredOptions
        $requiredOptions = self::toRequiredOptions($modelConfig);

        // Add input modalities if we have any inputs
        if (!empty($inputModalities)) {
            // Remove duplicates
            $inputModalities = array_unique($inputModalities, SORT_REGULAR);
            $requiredOptions = self::includeInRequiredOptions(
                $requiredOptions,
                new RequiredOption(OptionEnum::inputModalities(), array_values($inputModalities))
            );
        }

        return new self($capabilities, $requiredOptions);
    }

    /**
     * Determines the input modality contributed by a message part, if any.
     *
     * @since 1.4.0
     *
     * @param MessagePart $part The message part to analyze.
     * @return ModalityEnum|null The input modality, or null if the part contributes none.
     */
    private static function inputModalityForPart(MessagePart $part): ?ModalityEnum
    {
        // Check for text input
        if ($part->getType()->isText()) {
            return ModalityEnum::text();
        }

        // Check for file inputs
        if ($part->getType()->isFile()) {
            $file = $part->getFile();

            if ($file !== null) {
                if ($file->isImage()) {
                    return ModalityEnum::image();
                } elseif ($file->isAudio()) {
                    return ModalityEnum::audio();
                } elseif ($file->isVideo()) {
                    return ModalityEnum::video();
                } elseif ($file->isDocument() || $file->isText()) {
                    return ModalityEnum::document();
                }
            }
        }

        return null;
    }

    /**
     * Converts ModelConfig to an array of RequiredOptions.
     *
     * @since 0.2.0
     *
     * @param ModelConfig $modelConfig The model configuration.
     * @return list<RequiredOption> The required options.
     */
    private static function toRequiredOptions(ModelConfig $modelConfig): array
    {
        $requiredOptions = [];

        // Map properties that have corresponding OptionEnum values
        if ($modelConfig->getOutputModalities() !== null) {
            $requiredOptions[] = new RequiredOption(
                OptionEnum::outputModalities(),
                $modelConfig->getOutputModalities()
            );
        }

        if ($modelConfig->getSystemInstruction() !== null) {
            $requiredOptions[] = new RequiredOption(
                OptionEnum::systemInstruction(),
                $modelConfig->getSystemInstruction()
            );
        }

        if ($modelConfig->getCandidateCount() !== null) {
            $requiredOptions[] = new RequiredOption(
                OptionEnum::candidateCount(),
                $modelConfig->getCandidateCount()
            );
        }

        if ($modelConfig->getMaxTokens() !== null) {
            $requiredOptions[] = new RequiredOption(
                OptionEnum::maxTokens(),
                $modelConfig->getMaxTokens()
            );
        }

        if ($modelConfig->getTemperature() !== null) {
            $requiredOptions[] = new RequiredOption(
                OptionEnum::temperature(),
                $modelConfig->getTemperature()
            );
        }

        if ($modelConfig->getTopP() !== null) {
            $requiredOptions[] = new RequiredOption(
                OptionEnum::topP(),
                $modelConfig->getTopP()
            );
        }

        if ($modelConfig->getTopK() !== null) {
            $requiredOptions[] = new RequiredOption(
                OptionEnum::topK(),
                $modelConfig->getTopK()
            );
        }

        if ($modelConfig->getOutputMimeType() !== null) {
            $requiredOptions[] = new RequiredOption(
                OptionEnum::outputMimeType(),
                $modelConfig->getOutputMimeType()
            );
        }

        if ($modelConfig->getOutputSchema() !== null) {
            $requiredOptions[] = new RequiredOption(
                OptionEnum::outputSchema(),
                $modelConfig->getOutputSchema()
            );
        }

        // Handle properties without OptionEnum values as custom options
        if ($modelConfig->getStopSequences() !== null) {
            $requiredOptions[] = new RequiredOption(OptionEnum::stopSequences(), $modelConfig->getStopSequences());
        }

        if ($modelConfig->getPresencePenalty() !== null) {
            $requiredOptions[] = new RequiredOption(OptionEnum::presencePenalty(), $modelConfig->getPresencePenalty());
        }

        if ($modelConfig->getFrequencyPenalty() !== null) {
            $requiredOptions[] = new RequiredOption(
                OptionEnum::frequencyPenalty(),
                $modelConfig->getFrequencyPenalty()
            );
        }

        if ($modelConfig->getLogprobs() !== null) {
            $requiredOptions[] = new RequiredOption(OptionEnum::logprobs(), $modelConfig->getLogprobs());
        }

        if ($modelConfig->getTopLogprobs() !== null) {
            $requiredOptions[] = new RequiredOption(OptionEnum::topLogprobs(), $modelConfig->getTopLogprobs());
        }

        if ($modelConfig->getFunctionDeclarations() !== null) {
            $requiredOptions[] = new RequiredOption(OptionEnum::functionDeclarations(), true);
        }

        if ($modelConfig->getWebSearch() !== null) {
            $requiredOptions[] = new RequiredOption(OptionEnum::webSearch(), true);
        }

        if ($modelConfig->getOutputFileType() !== null) {
            $requiredOptions[] = new RequiredOption(OptionEnum::outputFileType(), $modelConfig->getOutputFileType());
        }

        if ($modelConfig->getOutputMediaOrientation() !== null) {
            $requiredOptions[] = new RequiredOption(
                OptionEnum::outputMediaOrientation(),
                $modelConfig->getOutputMediaOrientation()
            );
        }

        if ($modelConfig->getOutputMediaAspectRatio() !== null) {
            $requiredOptions[] = new RequiredOption(
                OptionEnum::outputMediaAspectRatio(),
                $modelConfig->getOutputMediaAspectRatio()
            );
        }

        if ($modelConfig->getDimensions() !== null) {
            $requiredOptions[] = new RequiredOption(OptionEnum::dimensions(), $modelConfig->getDimensions());
        }

        // Add custom options as individual RequiredOptions
        foreach ($modelConfig->getCustomOptions() as $key => $value) {
            $requiredOptions[] = new RequiredOption(OptionEnum::customOptions(), [$key => $value]);
        }

        return $requiredOptions;
    }

    /**
     * Includes a RequiredOption in the array, ensuring no duplicates based on option name.
     *
     * @since 0.2.0
     *
     * @param list<RequiredOption> $requiredOptions The existing required options.
     * @param RequiredOption $newOption The new option to include.
     * @return list<RequiredOption> The updated required options array.
     */
    private static function includeInRequiredOptions(array $requiredOptions, RequiredOption $newOption): array
    {
        // Check if we already have this option name
        foreach ($requiredOptions as $index => $existingOption) {
            if ($existingOption->getName()->equals($newOption->getName())) {
                // Replace existing option with new one
                $requiredOptions[$index] = $newOption;
                return $requiredOptions;
            }
        }

        // Option not found, add it
        $requiredOptions[] = $newOption;
        return $requiredOptions;
    }

    /**
     * {@inheritDoc}
     *
     * @since 0.1.0
     */
    public static function getJsonSchema(): array
    {
        return [
            'type' => 'object',
            'properties' => [
                self::KEY_REQUIRED_CAPABILITIES => [
                    'type' => 'array',
                    'items' => [
                        'type' => 'string',
                        'enum' => CapabilityEnum::getValues(),
                    ],
                    'description' => 'The capabilities that the model must support.',
                ],
                self::KEY_REQUIRED_OPTIONS => [
                    'type' => 'array',
                    'items' => RequiredOption::getJsonSchema(),
                    'description' => 'The options that the model must support with specific values.',
                ],
            ],
            'required' => [self::KEY_REQUIRED_CAPABILITIES, self::KEY_REQUIRED_OPTIONS],
        ];
    }

    /**
     * {@inheritDoc}
     *
     * @since 0.1.0
     *
     * @return ModelRequirementsArrayShape
     */
    public function toArray(): array
    {
        return [
            self::KEY_REQUIRED_CAPABILITIES => array_map(
                static fn(CapabilityEnum $capability): string => $capability->value,
                $this->requiredCapabilities
            ),
            self::KEY_REQUIRED_OPTIONS => array_map(
                static fn(RequiredOption $option): array => $option->toArray(),
                $this->requiredOptions
            ),
        ];
    }

    /**
     * {@inheritDoc}
     *
     * @since 0.1.0
     */
    public static function fromArray(array $array): self
    {
        static::validateFromArrayData($array, [self::KEY_REQUIRED_CAPABILITIES, self::KEY_REQUIRED_OPTIONS]);

        return new self(
            array_map(
                static fn(string $capability): CapabilityEnum => CapabilityEnum::from($capability),
                $array[self::KEY_REQUIRED_CAPABILITIES]
            ),
            array_map(
                static fn(array $optionData): RequiredOption => RequiredOption::fromArray($optionData),
                $array[self::KEY_REQUIRED_OPTIONS]
            )
        );
    }
}