<?php

declare( strict_types=1 );

namespace Fueled\AiProviderForOllama\Metadata;

use Fueled\AiProviderForOllama\Provider\OllamaProvider;
use WordPress\AiClient\Files\Enums\FileTypeEnum;
use WordPress\AiClient\Messages\Enums\ModalityEnum;
use WordPress\AiClient\Providers\ApiBasedImplementation\AbstractApiBasedModelMetadataDirectory;
use WordPress\AiClient\Providers\Http\DTO\Request;
use WordPress\AiClient\Providers\Http\DTO\RequestOptions;
use WordPress\AiClient\Providers\Http\Enums\HttpMethodEnum;
use WordPress\AiClient\Providers\Http\Exception\ResponseException;
use WordPress\AiClient\Providers\Http\Util\ResponseUtil;
use WordPress\AiClient\Providers\Models\DTO\ModelMetadata;
use WordPress\AiClient\Providers\Models\DTO\SupportedOption;
use WordPress\AiClient\Providers\Models\EmbeddingGeneration\Contracts\EmbeddingGenerationModelInterface;
use WordPress\AiClient\Providers\Models\Enums\CapabilityEnum;
use WordPress\AiClient\Providers\Models\Enums\OptionEnum;

/**
 * Class for the Ollama model metadata directory.
 *
 * @since 1.0.0
 *
 * @phpstan-type TagsEntryData array{
 *     name?: string,
 *     digest?: string,
 *     capabilities?: list<string>,
 *     details?: array{families?: list<string>|null}
 * }
 * @phpstan-type ShowResponseData array{
 *     capabilities?: list<string>,
 *     details?: array{families?: list<string>|null},
 *     voices?: list<string>|null
 * }
 * @phpstan-type ModelDetails array{capabilities: list<string>, families: list<string>, voices?: list<string>}
 */
class OllamaModelMetadataDirectory extends AbstractApiBasedModelMetadataDirectory {

	/**
	 * Default timeout for model discovery requests, in seconds.
	 *
	 * @since 1.3.0
	 *
	 * @var float
	 */
	private const DEFAULT_DISCOVERY_REQUEST_TIMEOUT = 10.0;

	/**
	 * Default connection timeout for model discovery requests, in seconds.
	 *
	 * @since 1.3.0
	 *
	 * @var float
	 */
	private const DEFAULT_DISCOVERY_CONNECT_TIMEOUT = 3.0;

	/**
	 * The model entries from /api/tags, once fetched.
	 *
	 * @since 1.3.0
	 *
	 * @var list<TagsEntryData>|null
	 */
	private ?array $model_tags = null;

	/**
	 * Lists the models the Ollama host offers, as returned by /api/tags.
	 *
	 * @since 1.3.0
	 *
	 * @return list<TagsEntryData> The raw model entries.
	 * @throws \WordPress\AiClient\Providers\Http\Exception\ResponseException If the host is unreachable or the response
	 *                                                                       is not a model listing.
	 */
	public function listModelTags(): array {
		if ( null !== $this->model_tags ) {
			return $this->model_tags;
		}

		$request  = $this->createRequest( HttpMethodEnum::GET(), 'api/tags' );
		$request  = $this->getRequestAuthentication()->authenticateRequest( $request );
		$response = $this->getHttpTransporter()->send( $request );

		ResponseUtil::throwIfNotSuccessful( $response );

		$tags_data = $response->getData();
		if ( ! isset( $tags_data['models'] ) || ! is_array( $tags_data['models'] ) ) {
			throw ResponseException::fromMissingData( 'Ollama', 'models' );
		}

		/** @var list<TagsEntryData> $model_tags */
		$model_tags       = array_values( array_filter( $tags_data['models'], 'is_array' ) );
		$this->model_tags = $model_tags;

		return $this->model_tags;
	}

	/**
	 * {@inheritDoc}
	 *
	 * @since 1.3.0
	 */
	public function invalidateCaches(): void {
		$this->model_tags = null;
		parent::invalidateCaches();
	}

	/**
	 * {@inheritDoc}
	 *
	 * @since 1.0.0
	 */
	protected function sendListModelsRequest(): array {
		$details_cache   = OllamaModelDetailsCache::load( OllamaProvider::url( '' ) );
		$digests_in_use  = array();
		$models_map      = array();
		$decision_models = array();

		foreach ( $this->listModelTags() as $model_entry ) {
			if ( ! isset( $model_entry['name'] ) || ! is_string( $model_entry['name'] ) || '' === $model_entry['name'] ) {
				continue;
			}

			$model_name = $model_entry['name'];
			$details    = $this->resolveModelDetails( $model_name, $model_entry, $details_cache, $digests_in_use );
			$metadata   = $this->buildModelMetadata( $model_name, $details );
			if ( null === $metadata ) {
				continue;
			}

			if ( null !== $details && in_array( 'decision', $details['capabilities'], true ) ) {
				$decision_models[ $model_name ] = $metadata;
				continue;
			}

			$models_map[ $model_name ] = $metadata;
		}

		$details_cache->save( $digests_in_use );

		ksort( $models_map );
		ksort( $decision_models );

		// Put Decision models last so they are not picked by default.
		return $models_map + $decision_models;
	}

	/**
	 * Lists the models that support decisions via the /v1/systemone endpoint.
	 *
	 * @since 1.3.0
	 *
	 * @return list<string> The model IDs, sorted.
	 * @throws \WordPress\AiClient\Providers\Http\Exception\ResponseException If the host is unreachable or the response
	 *                                                                       is not a model listing.
	 */
	public function listDecisionModelIds(): array {
		$model_ids = array();

		foreach ( $this->listModelTags() as $model_entry ) {
			if ( ! isset( $model_entry['name'] ) || ! is_string( $model_entry['name'] ) || '' === $model_entry['name'] ) {
				continue;
			}

			if ( ! in_array( 'decision', $this->readStringList( $model_entry['capabilities'] ?? null ), true ) ) {
				continue;
			}

			$model_ids[] = $model_entry['name'];
		}

		sort( $model_ids );

		return $model_ids;
	}

	/**
	 * Resolves the capability details of a single model, at the lowest cost available.
	 *
	 * In order of preference: the tag entry itself, the cache, and finally a
	 * request to /api/show.
	 *
	 * @since 1.3.0
	 *
	 * @param string                  $model_name     The model name.
	 * @param TagsEntryData           $model_entry    The model's entry from /api/tags.
	 * @param \Fueled\AiProviderForOllama\Metadata\OllamaModelDetailsCache $details_cache The cache of details fetched for earlier listings.
	 * @param list<string>            $digests_in_use Digests resolved from the cache so far, appended to by reference.
	 * @return ModelDetails|null The model details, or null when they could not be determined.
	 */
	private function resolveModelDetails(
		string $model_name,
		array $model_entry,
		OllamaModelDetailsCache $details_cache,
		array &$digests_in_use
	): ?array {
		$families = $this->readStringList( $model_entry['details']['families'] ?? null );

		// Recent Ollama versions report capabilities in the tag listing, making the per-model request unnecessary.
		// Speech models are the exception, as their voices are only reported by /api/show.
		$capabilities    = $this->readStringList( $model_entry['capabilities'] ?? null );
		$is_speech_model = in_array( 'speech', $capabilities, true );
		if ( ! empty( $capabilities ) && ! $is_speech_model ) {
			return array(
				'capabilities' => $capabilities,
				'families'     => $families,
			);
		}

		// The digest identifies the model's content, so an entry stays valid until the model itself changes.
		$digest = isset( $model_entry['digest'] ) && is_string( $model_entry['digest'] ) ? $model_entry['digest'] : '';

		if ( '' !== $digest ) {
			$cached_details = $details_cache->get( $digest );
			if ( null !== $cached_details ) {
				$digests_in_use[] = $digest;

				return $cached_details;
			}
		}

		$show_data = $this->fetchModelDetails( $model_name );
		if ( null === $show_data ) {
			if ( $is_speech_model ) {
				return array(
					'capabilities' => $capabilities,
					'families'     => $families,
				);
			}

			return null;
		}

		$show_families     = $this->readStringList( $show_data['details']['families'] ?? null );
		$show_capabilities = $this->readStringList( $show_data['capabilities'] ?? null );
		$details           = array(
			'capabilities' => empty( $show_capabilities ) ? $capabilities : $show_capabilities,
			'families'     => empty( $show_families ) ? $families : $show_families,
		);

		if ( in_array( 'speech', $details['capabilities'], true ) ) {
			$details['voices'] = $this->readStringList( $show_data['voices'] ?? null );
		}

		if ( '' !== $digest ) {
			$details_cache->set( $digest, $details );
			$digests_in_use[] = $digest;
		}

		return $details;
	}

	/**
	 * Builds a ModelMetadata object for a single model, or returns null if the model should be skipped.
	 *
	 * @since 1.0.0
	 *
	 * @param string $model_name The model name.
	 * @param ModelDetails|null $details The model's capability details, or null when they are unknown.
	 * @return \WordPress\AiClient\Providers\Models\DTO\ModelMetadata|null The model metadata, or null if the model should be excluded.
	 */
	private function buildModelMetadata( string $model_name, ?array $details ): ?ModelMetadata {
		$model_capabilities = null !== $details ? $details['capabilities'] : array();
		$model_families     = null !== $details ? $details['families'] : array();
		$model_voices       = null !== $details ? $details['voices'] ?? array() : array();

		// Speech models reject chat, completion, and embedding requests, so they only offer text to speech.
		if ( in_array( 'speech', $model_capabilities, true ) ) {
			return $this->buildTextToSpeechModelMetadata( $model_name, $model_voices );
		}

		$is_image_generation_model = in_array( 'image', $model_capabilities, true );

		$is_embedding_model = in_array( 'embedding', $model_capabilities, true )
			&& ! in_array( 'completion', $model_capabilities, true );

		if ( $is_embedding_model && ! $is_image_generation_model ) {
			// The embedding contracts are unreleased in some SDK versions.
			if ( ! interface_exists( EmbeddingGenerationModelInterface::class ) ) {
				return null;
			}

			return $this->buildEmbeddingModelMetadata( $model_name );
		}

		// Skip other non-completion models, but keep image-generation models which may not report "completion".
		if (
			! empty( $model_capabilities ) &&
			! in_array( 'completion', $model_capabilities, true ) &&
			! $is_image_generation_model
		) {
			return null;
		}

		// Check for vision support via the capabilities array or the model families.
		$has_vision = in_array( 'vision', $model_capabilities, true )
			|| in_array( 'clip', $model_families, true );

		$has_tools = in_array( 'tools', $model_capabilities, true );

		if ( $has_vision ) {
			$input_modalities_option = new SupportedOption(
				OptionEnum::inputModalities(),
				array(
					array( ModalityEnum::text() ),
					array( ModalityEnum::text(), ModalityEnum::image() ),
				)
			);
		} else {
			$input_modalities_option = new SupportedOption(
				OptionEnum::inputModalities(),
				array( array( ModalityEnum::text() ) )
			);
		}

		if ( $is_image_generation_model ) {
			return new ModelMetadata(
				$model_name,
				$model_name,
				array(
					CapabilityEnum::imageGeneration(),
				),
				array(
					new SupportedOption( OptionEnum::inputModalities(), array( array( ModalityEnum::text() ) ) ),
					new SupportedOption( OptionEnum::outputModalities(), array( array( ModalityEnum::image() ) ) ),
					new SupportedOption( OptionEnum::candidateCount() ),
					new SupportedOption( OptionEnum::outputMimeType(), array( 'image/png' ) ),
					new SupportedOption( OptionEnum::outputFileType(), array( FileTypeEnum::inline() ) ),
					new SupportedOption( OptionEnum::customOptions() ),
				)
			);
		}

		$options = array(
			new SupportedOption( OptionEnum::systemInstruction() ),
			new SupportedOption( OptionEnum::candidateCount() ),
			new SupportedOption( OptionEnum::maxTokens() ),
			new SupportedOption( OptionEnum::temperature() ),
			new SupportedOption( OptionEnum::topP() ),
			new SupportedOption( OptionEnum::topK() ),
			new SupportedOption( OptionEnum::stopSequences() ),
			new SupportedOption( OptionEnum::frequencyPenalty() ),
			new SupportedOption( OptionEnum::presencePenalty() ),
			new SupportedOption( OptionEnum::outputMimeType(), array( 'text/plain', 'application/json' ) ),
			new SupportedOption( OptionEnum::outputSchema() ),
			new SupportedOption( OptionEnum::customOptions() ),
			new SupportedOption( OptionEnum::outputModalities(), array( array( ModalityEnum::text() ) ) ),
			$input_modalities_option,
		);

		if ( $has_tools ) {
			$options[] = new SupportedOption( OptionEnum::functionDeclarations() );
		}

		return new ModelMetadata(
			$model_name,
			$model_name,
			array(
				CapabilityEnum::textGeneration(),
				CapabilityEnum::chatHistory(),
			),
			$options
		);
	}

	/**
	 * Builds embedding-generation metadata for a model.
	 *
	 * @since 1.2.0
	 *
	 * @param string $model_name The model name.
	 * @return \WordPress\AiClient\Providers\Models\DTO\ModelMetadata The embedding model metadata.
	 */
	private function buildEmbeddingModelMetadata( string $model_name ): ModelMetadata {
		return new ModelMetadata(
			$model_name,
			$model_name,
			array(
				CapabilityEnum::embeddingGeneration(),
			),
			array(
				new SupportedOption( OptionEnum::inputModalities(), array( array( ModalityEnum::text() ) ) ),
				new SupportedOption( OptionEnum::dimensions() ),
				new SupportedOption( OptionEnum::customOptions() ),
			)
		);
	}

	/**
	 * Builds text-to-speech metadata for a model.
	 *
	 * Ollama's speech endpoint only produces WAV (or raw PCM, which has no
	 * MIME type a browser can play), so WAV is the only output offered.
	 *
	 * @since 1.4.0
	 *
	 * @param string       $model_name The model name.
	 * @param list<string> $voices     The voices the model offers, or an empty list when they are unknown.
	 * @return \WordPress\AiClient\Providers\Models\DTO\ModelMetadata The text-to-speech model metadata.
	 */
	private function buildTextToSpeechModelMetadata( string $model_name, array $voices ): ModelMetadata {
		return new ModelMetadata(
			$model_name,
			$model_name,
			array(
				CapabilityEnum::textToSpeechConversion(),
			),
			array(
				new SupportedOption( OptionEnum::inputModalities(), array( array( ModalityEnum::text() ) ) ),
				new SupportedOption( OptionEnum::outputModalities(), array( array( ModalityEnum::audio() ) ) ),
				new SupportedOption( OptionEnum::outputMimeType(), array( 'audio/wav' ) ),
				new SupportedOption( OptionEnum::outputFileType(), array( FileTypeEnum::inline() ) ),
				// When the voices are unknown, any voice is accepted and Ollama validates it.
				new SupportedOption( OptionEnum::outputSpeechVoice(), empty( $voices ) ? null : $voices ),
				new SupportedOption( OptionEnum::customOptions() ),
			)
		);
	}

	/**
	 * Reads a list of strings out of an API payload.
	 *
	 * Ollama omits these keys on some versions and sends null for others, so
	 * anything that is not a list of strings is read as "none given".
	 *
	 * @since 1.3.0
	 *
	 * @param mixed $value The raw value.
	 * @return list<string> The strings it contained, if any.
	 */
	private function readStringList( $value ): array {
		if ( ! is_array( $value ) ) {
			return array();
		}

		return array_values( array_filter( $value, 'is_string' ) );
	}

	/**
	 * Fetches model details from the Ollama /api/show endpoint.
	 *
	 * Returns null if the request fails, in which case the caller falls back
	 * to default text-generation capabilities for the model.
	 *
	 * @since 1.0.0
	 *
	 * @param string $model_name The model name.
	 * @return ShowResponseData|null The response data, or null on failure.
	 */
	private function fetchModelDetails( string $model_name ): ?array {
		try {
			$request  = $this->createRequest(
				HttpMethodEnum::POST(),
				'api/show',
				array( 'Content-Type' => 'application/json' ),
				array(
					'model' => $model_name,
					'name'  => $model_name,
				)
			);
			$request  = $this->getRequestAuthentication()->authenticateRequest( $request );
			$response = $this->getHttpTransporter()->send( $request );

			ResponseUtil::throwIfNotSuccessful( $response );

			/** @var ShowResponseData $data */
			$data = $response->getData();
			return $data;
		} catch ( \Throwable $e ) {
			return null;
		}
	}

	/**
	 * Creates a request object for the Ollama API.
	 *
	 * @since 1.0.0
	 *
	 * @param \WordPress\AiClient\Providers\Http\Enums\HttpMethodEnum                     $method  The HTTP method.
	 * @param string                             $path    The API endpoint path, relative to the base URI.
	 * @param array<string, string|list<string>> $headers The request headers.
	 * @param string|array<string, mixed>|null   $data    The request data.
	 * @return \WordPress\AiClient\Providers\Http\DTO\Request The request object.
	 */
	private function createRequest( HttpMethodEnum $method, string $path, array $headers = array(), $data = null ): Request {
		return new Request(
			$method,
			OllamaProvider::url( $path ),
			$headers,
			$data,
			$this->discoveryRequestOptions()
		);
	}

	/**
	 * Builds the request options used for model discovery.
	 *
	 * @since 1.3.0
	 *
	 * @return \WordPress\AiClient\Providers\Http\DTO\RequestOptions The prepared request options.
	 */
	private function discoveryRequestOptions(): RequestOptions {
		$request_timeout = self::DEFAULT_DISCOVERY_REQUEST_TIMEOUT;
		$connect_timeout = self::DEFAULT_DISCOVERY_CONNECT_TIMEOUT;

		if ( function_exists( 'apply_filters' ) ) {
			/**
			 * Filters the request timeout for Ollama model discovery requests.
			 *
			 * Applies to the `/api/tags` and `/api/show` requests behind the connection
			 * check and the model list, not to text, image, or embedding generation.
			 *
			 * @since 1.3.0
			 *
			 * @param float $request_timeout The request timeout in seconds.
			 */
			$request_timeout = (float) apply_filters( 'ai_provider_for_ollama_discovery_request_timeout', $request_timeout );

			/**
			 * Filters the connection timeout for Ollama model discovery requests.
			 *
			 * @since 1.3.0
			 *
			 * @param float $connect_timeout The connection timeout in seconds.
			 */
			$connect_timeout = (float) apply_filters( 'ai_provider_for_ollama_discovery_connect_timeout', $connect_timeout );
		}

		$request_options = new RequestOptions();
		$request_options->setTimeout( $request_timeout );
		$request_options->setConnectTimeout( $connect_timeout );

		return $request_options;
	}
}