import { RagMilvusRuntimeConfig } from '../config'
import {
  createMilvusSdkClient,
  MilvusVectorStoreClient,
} from '../milvus-sdk.client'
import { cosineSimilarity } from '../../utils/math'
import {
  RagVectorSearchInput,
  RagVectorSearchMatch,
  RagVectorStore,
  RagVectorChunkPayload,
} from '../types'
import { mapLegacySourceTypeToRetrievalSourceType } from '../../rag.types'
import type { RagRetrievalSourceType } from '../../rag.types'

/**
 * Milvus 向量存储适配器（教学骨架）。
 *
 * 说明：
 * - 本轮仅提供 mock 模式，便于先完成契约回归；
 * - `sdk` 模式留给后续步骤接入真实 Milvus Client。
 */
export class MilvusRagVectorStoreAdapter implements RagVectorStore {
  readonly backend = 'milvus' as const
  private readonly mockStore = new Map<string, RagVectorChunkPayload>()
  private sdkClient: MilvusVectorStoreClient | null

  constructor(
    private readonly config: RagMilvusRuntimeConfig,
    sdkClient?: MilvusVectorStoreClient,
  ) {
    this.sdkClient = sdkClient ?? null
  }

  async upsertChunks(chunks: RagVectorChunkPayload[]): Promise<void> {
    if (this.config.mode === 'sdk') {
      await this.getSdkClient().upsertChunks(chunks)
      return
    }

    for (const chunk of chunks) {
      this.mockStore.set(chunk.id, chunk)
    }
  }

  async deleteChunksByDocument(documentId: string): Promise<void> {
    if (this.config.mode === 'sdk') {
      await this.getSdkClient().deleteChunksByDocument(documentId)
      return
    }

    for (const [id, chunk] of this.mockStore.entries()) {
      if (chunk.documentId === documentId) {
        this.mockStore.delete(id)
      }
    }
  }

  async listDocumentIds(sourceType?: RagRetrievalSourceType): Promise<string[]> {
    if (this.config.mode === 'sdk') {
      return this.getSdkClient().listDocumentIds(sourceType)
    }

    const documentIds = Array.from(this.mockStore.values())
      .filter((chunk) =>
        sourceType
          ? mapLegacySourceTypeToRetrievalSourceType(chunk.sourceType) === sourceType
          : true,
      )
      .map((chunk) => chunk.documentId)

    return Array.from(new Set(documentIds))
  }

  async search(input: RagVectorSearchInput): Promise<RagVectorSearchMatch[]> {
    if (this.config.mode === 'sdk') {
      return this.getSdkClient().search(input)
    }

    const candidates = Array.from(this.mockStore.values())
      .filter((chunk) =>
        input.sourceType ? chunk.sourceType === input.sourceType : true,
      )
      .filter((chunk) =>
        input.sourceTypes?.length
          ? input.sourceTypes.includes(mapLegacySourceTypeToRetrievalSourceType(chunk.sourceType))
          : true,
      )
      .filter((chunk) =>
        input.sourceScope ? chunk.sourceScope === input.sourceScope : true,
      )
      .filter((chunk) =>
        input.knowledgeDomains?.length
          ? input.knowledgeDomains.includes(
              chunk.metadataJson?.knowledgeDomain as (typeof input.knowledgeDomains)[number],
            )
          : true,
      )
      .filter((chunk) =>
        input.documentIds?.length
          ? input.documentIds.includes(chunk.documentId)
          : true,
      )

    return candidates
      .map((chunk) => ({
        ...chunk,
        score: Number(cosineSimilarity(input.queryVector, chunk.embedding).toFixed(6)),
      }))
      .sort((left, right) => right.score - left.score)
      .slice(0, Math.max(Math.floor(input.limit), 0))
  }

  private getSdkClient(): MilvusVectorStoreClient {
    if (!this.sdkClient) {
      this.sdkClient = createMilvusSdkClient(this.config)
    }

    return this.sdkClient
  }
}
