//
// Copyright (c) 2023-2026, Manticore Software LTD (https://manticoresearch.com)
// All rights reserved
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License. You should have
// received a copy of the GPL license along with this program; if you
// did not, you can find it at http://www.gnu.org/
//

#pragma once

#include "sphinxstd.h"
#include "sphinxdefs.h"
#include "columnarmisc.h"
#include "indexsettings.h"
#include "secondaryindex.h"

#include <atomic>
#include <functional>

class CSphQueryContext;
class QueryProfile_c;

class TableEmbeddings_c
{
public:
	bool	Load ( const CSphString & sAttr, const knn::ModelSettings_t & tSettings, CSphString & sError );
	knn::TextToEmbeddings_i * GetModel ( const CSphString & sAttr ) const;

private:
	SmallStringHash_T<std::unique_ptr<knn::TextToEmbeddings_i>> m_hModels;
};

class EmbeddingsSrc_c
{
public:
	struct Row_t
	{
		CSphVector<char>	m_dSrc;
		bool				m_bDefault = false;
		void SwapData ( Row_t & rhs );
	};

			EmbeddingsSrc_c ( int iAttrs );

	void	Add ( int iAttr, CSphVector<char> & dSrc, bool bDefault );
	void	SwapRows ( RowID_t tDstID, RowID_t tSrcID );
	void	DropTail ( RowID_t tTailID );
	const VecTraits_T<char> Get ( RowID_t tRowID, int iAttr ) const;
	bool	Has ( RowID_t tRowID, int iAttr ) const;
	bool	IsDefault ( RowID_t tRowID, int iAttr ) const;

private:
	CSphVector<CSphVector<Row_t>> m_dRows;
};

bool							IsKnnDist ( const CSphString & sExpr );
const char *					GetKnnDistAttrName();
const char *					GetKnnDistRescoreAttrName();
void							SetupKNNLimit ( CSphQuery & tQuery );

ISphExpr *						CreateExpr_KNNDist ( const CSphVector<float> & dAnchor, const CSphColumnInfo & tAttr );
void							NormalizeVec ( VecTraits_T<float> & dData );

void							AddKNNSettings ( StringBuilder_c & sRes, const CSphColumnInfo & tAttr );
void							ReadKNNJson ( bson::Bson_c tRoot, knn::IndexSettings_t & tIS, knn::ModelSettings_t & tMS, CSphString & sKNNFrom );
CSphString						FormatKNNConfigStr ( const CSphVector<NamedKNNSettings_t> & dAttrs );
bool							ParseKNNConfigStr ( const CSphString & sStr, CSphVector<NamedKNNSettings_t> & dParsed, CSphString & sError );
void							FormatKNNSettings ( JsonEscapedBuilder & tOut, const knn::IndexSettings_t & tIndexSettings, const knn::ModelSettings_t & tModelSettings, const CSphString & sKNNFrom );

bool							Str2HNSWSimilarity ( const CSphString & sSimilarity, knn::HNSWSimilarity_e & eSimilarity, CSphString * pError = nullptr );
bool							Str2Quantization ( const CSphString & sQuantization, knn::Quantization_e & eQuantization, CSphString * pError = nullptr );

int								GetDefaultKNNParallelBuild ( int iThreads );
int								GetKNNBuildConcurrency ( int64_t iTotalRows );
CSphString						ConcatKNNBuildErrors ( const VecTraits_T<CSphString> & dErrors );
void							RecordKNNBuildError ( CSphString & sSlot, const knn::BuildContext_t & tBuildCtx, const char * szFallbackFmt, ... ) __attribute__ ( ( format ( printf, 3, 4 ) ) );

std::unique_ptr<knn::Builder_i> BuildCreateKNN ( const ISphSchema & tSchema, int64_t iNumElements, CSphVector<std::pair<PlainOrColumnar_t,int>> & dAttrs, const CSphString & sTmpFilename, CSphString & sError );
void							BuildTrainKNN ( RowID_t tRowIDSrc, RowID_t tRowIDDst, const CSphRowitem * pRow, const BYTE * pPool, CSphVector<ScopedTypedIterator_t> & dIterators, const VecTraits_T<PlainOrColumnar_t> & dAttrs, knn::Builder_i & tBuilder );
bool							BuildStoreKNN ( RowID_t tRowIDSrc, RowID_t tRowIDDst, const CSphRowitem * pRow, const BYTE * pPool, CSphVector<ScopedTypedIterator_t> & dIterators, const VecTraits_T<PlainOrColumnar_t> & dAttrs, knn::Builder_i & tBuilder, knn::BuildContext_t & tBuildCtx );
bool							BuildStoreKNNParallelDiskIndex ( const CSphIndex & tIndex, knn::Builder_i & tBuilder, const VecTraits_T<PlainOrColumnar_t> & dAttrs, int64_t iTotalRows, CSphString & sError );

using KNNStoreWorkerFn_t = std::function<void ( int iWorkerIdx, int64_t iStart, int64_t iEnd, CSphString & sWorkerError, std::atomic<bool> & bStop )>;
bool							RunParallelKNNStore ( int64_t iTotalUnits, CSphString & sError, KNNStoreWorkerFn_t fnWorker );

std::pair<RowidIterator_i *, bool> CreateKNNIterator ( knn::KNN_i * pKNN, const CSphQuery & tQuery, const ISphSchema & tIndexSchema, const ISphSchema & tSorterSchema, knn::KNNFilter_i * pFilter, knn::HNSWTerminationPolicy_e ePolicy, QueryProfile_c * pProfile, CSphString & sError );
RowIteratorsWithEstimates_t		CreateKNNIterators ( knn::KNN_i * pKNN, const CSphQuery & tQuery, const ISphSchema & tIndexSchema, const ISphSchema & tSorterSchema, knn::KNNFilter_i * pFilter, knn::HNSWTerminationPolicy_e ePolicy, QueryProfile_c * pProfile, bool & bError, CSphString & sError );
std::unique_ptr<knn::KNNFilter_i> CreateKNNPrefilter ( const CSphQueryContext & tCtx, const CSphRowitem * pAttrPool, int iStride, int iDynamicSize, int64_t iFilterCount );

ISphMatchSorter *				CreateKNNRescoreSorter ( ISphMatchSorter * pSorter, const KnnSearchSettings_t & tSettings, ESphSortFunc eMatchFunc );

const char *					GetAPITimeoutErrorMsg();
bool							ValidateEmbeddingsAPITimeout ( const CSphString & sValue, int & iTimeout, CSphString & sError );
