import { TokenManager, PluginEndpointDiscovery } from '@backstage/backend-common';
import { CatalogApi } from '@backstage/catalog-client';
import { LoggerService, AuthService, DiscoveryService } from '@backstage/backend-plugin-api';
import { RoadieVectorStore, AugmentationIndexer, EmbeddingsSource, EmbeddingDoc, EntityFilterShape, RetrievalPipeline, RetrievalRouter, AugmentationRetriever, AugmentationPostProcessor } from '@roadiehq/rag-ai-node';
import { Entity } from '@backstage/catalog-model';
import { Embeddings } from '@langchain/core/embeddings';
import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter';

type AugmentationOptions = {
    chunkSize?: number;
    chunkOverlap?: number;
    concurrencyLimit?: number;
};
interface RoadieEmbeddingsConfig {
    logger: LoggerService;
    tokenManager?: TokenManager;
    auth?: AuthService;
    vectorStore: RoadieVectorStore;
    catalogApi: CatalogApi;
    discovery: PluginEndpointDiscovery;
    augmentationOptions?: AugmentationOptions;
}
type TechDocsDocument = {
    text: string;
    entity: Entity;
    title: string;
    location: string;
};

declare class DefaultVectorAugmentationIndexer implements AugmentationIndexer {
    private readonly _vectorStore;
    private readonly catalogApi;
    private readonly logger;
    private readonly auth;
    private readonly discovery;
    private readonly augmentationOptions?;
    protected constructor({ vectorStore, catalogApi, logger, auth, tokenManager, embeddings, discovery, augmentationOptions, }: {
        vectorStore: RoadieVectorStore;
        catalogApi: CatalogApi;
        logger: LoggerService;
        auth?: AuthService;
        tokenManager?: TokenManager;
        embeddings: Embeddings;
        discovery: DiscoveryService;
        augmentationOptions?: AugmentationOptions;
    });
    get vectorStore(): RoadieVectorStore;
    /**
     * Returns the splitter object. Default implementation is using a naive RecursiveCharacterTextSplitter
     * which is likely not the best candidate for structured data splitting.
     *
     * It is recommended that this method is overwritten with more applicable implementation
     *
     * @returns {RecursiveCharacterTextSplitter} The splitter object.
     */
    protected getSplitter(): RecursiveCharacterTextSplitter;
    protected constructCatalogEmbeddingDocuments(entities: Entity[], source: EmbeddingsSource): Promise<EmbeddingDoc[]>;
    protected constructTechDocsEmbeddingDocuments(documents: TechDocsDocument[], source: EmbeddingsSource): Promise<EmbeddingDoc[]>;
    protected getDocuments(source: EmbeddingsSource, filter?: EntityFilterShape): Promise<EmbeddingDoc[]>;
    createEmbeddings(source: EmbeddingsSource, filter: EntityFilterShape): Promise<number>;
    deleteEmbeddings(source: EmbeddingsSource, filter: EntityFilterShape): Promise<void>;
}

declare class DefaultRetrievalPipeline implements RetrievalPipeline {
    private readonly routers;
    private readonly retrievers;
    private readonly postProcessors;
    constructor({ routers, retrievers, postProcessors, }: {
        routers?: RetrievalRouter[];
        retrievers?: AugmentationRetriever[];
        postProcessors?: AugmentationPostProcessor[];
    });
    retrieveAugmentationContext(query: string, source: EmbeddingsSource, filter?: EntityFilterShape): Promise<EmbeddingDoc[]>;
}

declare class CombiningPostProcessor implements AugmentationPostProcessor {
    process(_query: string, _source: EmbeddingsSource, embeddingDocs: Map<string, EmbeddingDoc[]>): Promise<EmbeddingDoc[]>;
}

declare class VectorEmbeddingsRetriever implements AugmentationRetriever {
    private readonly logger;
    private readonly vectorStore;
    constructor({ vectorStore, logger, }: {
        vectorStore: RoadieVectorStore;
        logger: LoggerService;
    });
    get id(): string;
    retrieve(query: string, source: EmbeddingsSource, _filter?: EntityFilterShape): Promise<EmbeddingDoc[]>;
}

type SearchClientQuery = {
    term: string;
    source: EmbeddingsSource;
};
declare class SearchClient {
    private readonly discoveryApi;
    private readonly logger;
    private readonly auth;
    constructor(options: {
        discoveryApi: PluginEndpointDiscovery;
        logger: LoggerService;
        auth?: AuthService;
        tokenManager?: TokenManager;
    });
    query(query: SearchClientQuery): Promise<EmbeddingDoc[]>;
}

declare class SearchRetriever implements AugmentationRetriever {
    private readonly searchClient;
    private readonly logger;
    constructor({ discovery, logger, searchClient, auth, tokenManager, }: {
        discovery: PluginEndpointDiscovery;
        logger: LoggerService;
        searchClient?: SearchClient;
        auth?: AuthService;
        tokenManager?: TokenManager;
    });
    get id(): string;
    retrieve(query: string, source: EmbeddingsSource): Promise<EmbeddingDoc[]>;
}

declare class SourceBasedRetrievalRouter implements RetrievalRouter {
    private readonly logger;
    private readonly retrievers;
    constructor({ logger, retrievers, }: {
        logger: LoggerService;
        retrievers: Map<EmbeddingsSource, AugmentationRetriever[]>;
    });
    determineRetriever(_query: string, source: EmbeddingsSource): Promise<AugmentationRetriever[]>;
}

type DefaultRetrievalPipelineOptions = {
    vectorStore: RoadieVectorStore;
    logger: LoggerService;
    discovery: PluginEndpointDiscovery;
    auth?: AuthService;
    tokenManager?: TokenManager;
};
declare const createDefaultRetrievalPipeline: ({ vectorStore, discovery, logger, auth, tokenManager, }: DefaultRetrievalPipelineOptions) => DefaultRetrievalPipeline;

export { type AugmentationOptions, CombiningPostProcessor, DefaultRetrievalPipeline, type DefaultRetrievalPipelineOptions, DefaultVectorAugmentationIndexer, type RoadieEmbeddingsConfig, SearchRetriever, SourceBasedRetrievalRouter, VectorEmbeddingsRetriever, createDefaultRetrievalPipeline };