<?php

namespace App\Services;

use App\Models\EmbeddingEntry;
use Illuminate\Support\Collection;
use Illuminate\Support\Facades\DB;
use Illuminate\Support\Facades\Http;
use InvalidArgumentException;
use Laravel\Ai\Ai;
use Laravel\Ai\Embeddings;
use Laravel\Ai\Enums\Lab;
use RuntimeException;
use Throwable;

class EmbeddingWorkbench
{
    /**
     * @var Collection<int, string>|null
     */
    protected ?Collection $ollamaModelNames = null;

    public function providerName(): string
    {
        $provider = config('ai.default_for_embeddings', 'ollama');

        return $provider instanceof Lab ? $provider->value : (string) $provider;
    }

    public function modelName(?string $model = null): string
    {
        $model = trim((string) $model);

        if ($model !== '') {
            return $this->validateModel($model);
        }

        $default = Ai::embeddingProvider($this->providerName())->defaultEmbeddingsModel();
        $availableModels = $this->availableModelNames();

        if ($availableModels->contains($default)) {
            return $default;
        }

        $firstAvailableModel = $availableModels->first();

        if ($firstAvailableModel === null) {
            throw new InvalidArgumentException('No embedding models are available.');
        }

        return $firstAvailableModel;
    }

    public function dimensions(?string $model = null): int
    {
        $this->modelName($model);

        return Ai::embeddingProvider($this->providerName())->defaultEmbeddingsDimensions();
    }

    /**
     * @return array<int, array{name: string, dimensions: int, default: bool}>
     */
    public function modelOptions(): array
    {
        $default = Ai::embeddingProvider($this->providerName())->defaultEmbeddingsModel();

        return $this->availableModelNames()
            ->map(fn (string $model): array => [
                'name' => $model,
                'dimensions' => $this->dimensions($model),
                'default' => $model === $default,
            ])
            ->all();
    }

    /**
     * @param  list<string>  $phrases
     * @return Collection<int, EmbeddingEntry>
     */
    public function embed(array $phrases, ?string $model = null): Collection
    {
        $phrases = $this->uniquePhrases($phrases);

        if ($phrases === []) {
            return collect();
        }

        $provider = $this->providerName();
        $model = $this->modelName($model);
        $dimensions = $this->dimensions($model);
        $hashes = collect($phrases)->mapWithKeys(fn (string $phrase): array => [
            EmbeddingVector::hashPhrase($phrase) => $phrase,
        ]);

        $entries = EmbeddingEntry::query()
            ->where('provider', $provider)
            ->where('model', $model)
            ->where('embedding_dimensions', $dimensions)
            ->whereIn('content_hash', $hashes->keys())
            ->get()
            ->keyBy('content_hash');

        $missingPhrases = $hashes
            ->reject(fn (string $phrase, string $hash): bool => $entries->has($hash))
            ->values()
            ->all();

        if ($missingPhrases !== []) {
            $response = Embeddings::for($missingPhrases)
                ->dimensions($dimensions)
                ->cache()
                ->timeout(60)
                ->generate($provider, $model);

            foreach ($missingPhrases as $index => $phrase) {
                $embedding = array_map('floatval', $response->embeddings[$index] ?? []);

                $this->ensureExpectedDimensions($embedding, $dimensions);

                $entry = EmbeddingEntry::query()->updateOrCreate(
                    [
                        'content_hash' => EmbeddingVector::hashPhrase($phrase),
                        'provider' => $response->meta->provider,
                        'model' => $response->meta->model,
                        'embedding_dimensions' => $dimensions,
                    ],
                    [
                        'source_text' => $phrase,
                        'embedding' => $embedding,
                        'tokens' => $response->tokens,
                    ]
                );

                $entries->put($entry->content_hash, $entry);
            }
        }

        return collect($phrases)
            ->map(fn (string $phrase): ?EmbeddingEntry => $entries->get(EmbeddingVector::hashPhrase($phrase)))
            ->filter()
            ->values();
    }

    /**
     * @return Collection<int, array{entry: EmbeddingEntry, similarity: float, distance: float}>
     */
    public function search(string $query, int $limit = 10, float $minimumSimilarity = 0.3, ?string $model = null): Collection
    {
        $query = EmbeddingVector::cleanPhrase($query);

        if ($query === '') {
            return collect();
        }

        $limit = max(1, min(50, $limit));
        $minimumSimilarity = max(0.0, min(1.0, $minimumSimilarity));
        $provider = $this->providerName();
        $model = $this->modelName($model);
        $dimensions = $this->dimensions($model);
        $queryEmbedding = $this->generateEmbedding($query, $model);

        if ($this->canUseVectorQueries()) {
            return EmbeddingEntry::query()
                ->select('embedding_entries.*')
                ->selectVectorDistance('embedding', $queryEmbedding, as: 'distance')
                ->where('provider', $provider)
                ->where('model', $model)
                ->where('embedding_dimensions', $dimensions)
                ->whereVectorSimilarTo('embedding', $queryEmbedding, minSimilarity: $minimumSimilarity)
                ->limit($limit)
                ->get()
                ->map(fn (EmbeddingEntry $entry): array => [
                    'entry' => $entry,
                    'similarity' => round(1 - (float) $entry->distance, 6),
                    'distance' => round((float) $entry->distance, 6),
                ]);
        }

        return EmbeddingEntry::query()
            ->where('provider', $provider)
            ->where('model', $model)
            ->where('embedding_dimensions', $dimensions)
            ->get()
            ->map(function (EmbeddingEntry $entry) use ($queryEmbedding): array {
                $similarity = EmbeddingVector::cosineSimilarity($queryEmbedding, $entry->embedding);

                return [
                    'entry' => $entry,
                    'similarity' => round($similarity, 6),
                    'distance' => round(1 - $similarity, 6),
                ];
            })
            ->filter(fn (array $result): bool => $result['similarity'] >= $minimumSimilarity)
            ->sortByDesc('similarity')
            ->take($limit)
            ->values();
    }

    /**
     * @param  list<string>  $phrases
     * @return array{entries: Collection<int, EmbeddingEntry>, matrix: array<int, array{entry: EmbeddingEntry, scores: array<int, float>}>, closest_pair: array{first: EmbeddingEntry, second: EmbeddingEntry, similarity: float}|null}
     */
    public function compare(array $phrases, ?string $model = null): array
    {
        $entries = $this->embed($phrases, $model)->values();
        $matrix = [];
        $closestPair = null;

        foreach ($entries as $rowIndex => $entry) {
            $scores = [];

            foreach ($entries as $columnIndex => $comparedEntry) {
                $similarity = round(EmbeddingVector::cosineSimilarity($entry->embedding, $comparedEntry->embedding), 6);
                $scores[] = $similarity;

                if ($rowIndex < $columnIndex && (
                    $closestPair === null || $similarity > $closestPair['similarity']
                )) {
                    $closestPair = [
                        'first' => $entry,
                        'second' => $comparedEntry,
                        'similarity' => $similarity,
                    ];
                }
            }

            $matrix[] = [
                'entry' => $entry,
                'scores' => $scores,
            ];
        }

        return [
            'entries' => $entries,
            'matrix' => $matrix,
            'closest_pair' => $closestPair,
        ];
    }

    /**
     * @return array{id: int, source_text: string, provider: string, model: string, dimensions: int, tokens: int|null, vector_preview: array<int, float>, vector_json: string}
     */
    public function presentEntry(EmbeddingEntry $entry): array
    {
        return [
            'id' => $entry->id,
            'source_text' => $entry->source_text,
            'provider' => $entry->provider,
            'model' => $entry->model,
            'dimensions' => $entry->embedding_dimensions,
            'tokens' => $entry->tokens,
            'vector_preview' => EmbeddingVector::preview($entry->embedding),
            'vector_json' => json_encode(EmbeddingVector::rounded($entry->embedding), JSON_PRETTY_PRINT | JSON_THROW_ON_ERROR),
        ];
    }

    /**
     * @param  array<int, string>  $phrases
     * @return list<string>
     */
    protected function uniquePhrases(array $phrases): array
    {
        $unique = [];

        foreach ($phrases as $phrase) {
            $phrase = EmbeddingVector::cleanPhrase($phrase);

            if ($phrase === '') {
                continue;
            }

            $unique[EmbeddingVector::hashPhrase($phrase)] = $phrase;
        }

        return array_values($unique);
    }

    /**
     * @return array<int, float>
     */
    protected function generateEmbedding(string $text, ?string $model = null): array
    {
        $model = $this->modelName($model);
        $dimensions = $this->dimensions($model);
        $embedding = array_map(
            'floatval',
            Embeddings::for([$text])
                ->dimensions($dimensions)
                ->cache()
                ->timeout(60)
                ->generate($this->providerName(), $model)
                ->first()
        );

        $this->ensureExpectedDimensions($embedding, $dimensions);

        return $embedding;
    }

    /**
     * @param  array<int, float>  $embedding
     */
    protected function ensureExpectedDimensions(array $embedding, int $dimensions): void
    {
        if (count($embedding) !== $dimensions) {
            throw new RuntimeException('The embedding model returned '.count($embedding)." dimensions, but the database is configured for {$dimensions}.");
        }
    }

    protected function canUseVectorQueries(): bool
    {
        return DB::connection()->getDriverName() === 'pgsql';
    }

    /**
     * @return Collection<int, string>
     */
    protected function availableModelNames(): Collection
    {
        if ($this->providerName() === 'ollama') {
            $ollamaModels = $this->ollamaModelNames();

            if ($ollamaModels->isNotEmpty()) {
                return $ollamaModels;
            }
        }

        return $this->configuredModelNames();
    }

    /**
     * @return Collection<int, string>
     */
    protected function configuredModelNames(): Collection
    {
        $models = config('ai.providers.'.$this->providerName().'.models.embeddings.available', []);

        if (! is_array($models)) {
            $models = [];
        }

        return collect($models)
            ->push(Ai::embeddingProvider($this->providerName())->defaultEmbeddingsModel())
            ->map(fn (mixed $model): string => trim((string) $model))
            ->filter()
            ->unique()
            ->values();
    }

    /**
     * @return Collection<int, string>
     */
    protected function ollamaModelNames(): Collection
    {
        if ($this->ollamaModelNames !== null) {
            return $this->ollamaModelNames;
        }

        try {
            $models = Http::baseUrl($this->ollamaUrl())
                ->acceptJson()
                ->connectTimeout(2)
                ->timeout(5)
                ->get('api/tags')
                ->throw()
                ->json('models', []);
        } catch (Throwable) {
            return $this->ollamaModelNames = collect();
        }

        if (! is_array($models)) {
            return $this->ollamaModelNames = collect();
        }

        return $this->ollamaModelNames = collect($models)
            ->map(fn (mixed $model): string => is_array($model) ? trim((string) ($model['name'] ?? '')) : '')
            ->filter()
            ->unique()
            ->values();
    }

    protected function ollamaUrl(): string
    {
        return rtrim((string) config('ai.providers.ollama.url', 'http://localhost:11434'), '/');
    }

    protected function validateModel(string $model): string
    {
        $model = trim($model);

        if ($model === '') {
            throw new InvalidArgumentException('Choose an embedding model.');
        }

        $availableModels = $this->availableModelNames();

        if ($availableModels->doesntContain($model)) {
            throw new InvalidArgumentException("The embedding model [{$model}] is not available.");
        }

        return $model;
    }
}