/**
 * NAC 公链 - 向量 Embedding 检索模块
 * Vector Embedding Retrieval Module
 *
 * 功能：
 * 1. 将规则文本转换为向量表示（使用内置 LLM API）
 * 2. 计算查询向量与规则向量的余弦相似度
 * 3. 返回语义最相关的规则（替代正则关键词匹配）
 * 4. 支持 MongoDB Atlas Vector Search（生产环境）
 *    和内存向量检索（降级模式）
 *
 * 架构：
 * - 优先使用 MongoDB Atlas Vector Search（如果可用）
 * - 降级到内存向量检索（TF-IDF + 余弦相似度）
 * - 最终降级到现有的正则关键词检索
 */

import { MongoClient, Collection, Document } from "mongodb";

// ─── 类型定义 ─────────────────────────────────────────────────────

export interface EmbeddingVector {
  ruleId: string;
  vector: number[];
  text: string;
  createdAt: Date;
}

export interface SemanticSearchResult {
  ruleId: string;
  ruleName: string;
  jurisdiction: string;
  assetClass: string;
  ruleType: string;
  content: string;
  legalBasis?: string;
  ownershipRequirements?: Record<string, unknown>;
  tradingRequirements?: Record<string, unknown>;
  score: number;
  similarityScore: number;
  sourceUrl?: string;
  tags?: string[];
  complianceLevel?: string;
}

// ─── TF-IDF 向量化（内存模式，无需外部 API）─────────────────────

/**
 * 构建 TF-IDF 词汇表
 * 支持中英文混合文本
 */
class TFIDFVectorizer {
  private vocabulary: Map<string, number> = new Map();
  private idf: Map<string, number> = new Map();
  private documents: string[][] = [];

  /**
   * 分词（支持中英文）
   */
  tokenize(text: string): string[] {
    const normalized = text.toLowerCase()
      .replace(/[^\u4e00-\u9fa5a-z0-9\s]/g, " ")
      .replace(/\s+/g, " ")
      .trim();

    const tokens: string[] = [];

    // 英文单词（空格分割）
    const englishWords = normalized.match(/[a-z][a-z0-9]*/g) || [];
    tokens.push(...englishWords.filter(w => w.length > 2));

    // 中文字符（2-4字 n-gram）
    const chineseText = normalized.replace(/[a-z0-9\s]/g, "");
    for (let i = 0; i < chineseText.length - 1; i++) {
      // 双字词
      tokens.push(chineseText.slice(i, i + 2));
      // 三字词
      if (i < chineseText.length - 2) {
        tokens.push(chineseText.slice(i, i + 3));
      }
    }

    return tokens;
  }

  /**
   * 拟合语料库，建立词汇表和 IDF
   */
  fit(documents: string[]): void {
    this.documents = documents.map(doc => this.tokenize(doc));

    // 建立词汇表
    const allTokensSet = new Set<string>();
    for (const tokens of this.documents) {
      for (const token of tokens) {
        allTokensSet.add(token);
      }
    }
    const allTokens = Array.from(allTokensSet);

    let idx = 0;
    for (const token of allTokens) {
      this.vocabulary.set(token, idx++);
    }

    // 计算 IDF
    const N = this.documents.length;
    for (const token of allTokens) {
      const df = this.documents.filter(doc => doc.includes(token)).length;
      this.idf.set(token, Math.log((N + 1) / (df + 1)) + 1);
    }
  }

  /**
   * 将文本转换为 TF-IDF 向量
   */
  transform(text: string): number[] {
    const tokens = this.tokenize(text);
    const vector = new Array(this.vocabulary.size).fill(0);

    // 计算 TF
    const tf = new Map<string, number>();
    for (const token of tokens) {
      tf.set(token, (tf.get(token) || 0) + 1);
    }

    // 计算 TF-IDF
    for (const [token, count] of Array.from(tf.entries())) {
      const idx = this.vocabulary.get(token);
      if (idx !== undefined) {
        const tfScore = count / tokens.length;
        const idfScore = this.idf.get(token) || 1;
        vector[idx] = tfScore * idfScore;
      }
    }

    // L2 归一化
    const norm = Math.sqrt(vector.reduce((sum, v) => sum + v * v, 0));
    if (norm > 0) {
      return vector.map(v => v / norm);
    }
    return vector;
  }

  getVocabularySize(): number {
    return this.vocabulary.size;
  }
}

// ─── 余弦相似度计算 ───────────────────────────────────────────────

function cosineSimilarity(a: number[], b: number[]): number {
  if (a.length !== b.length) return 0;

  let dotProduct = 0;
  let normA = 0;
  let normB = 0;

  for (let i = 0; i < a.length; i++) {
    dotProduct += a[i] * b[i];
    normA += a[i] * a[i];
    normB += b[i] * b[i];
  }

  const denominator = Math.sqrt(normA) * Math.sqrt(normB);
  if (denominator === 0) return 0;

  return dotProduct / denominator;
}

// ─── 内存向量检索引擎 ─────────────────────────────────────────────

interface RuleVector {
  doc: Record<string, unknown>;
  vector: number[];
  text: string;
}

class InMemoryVectorSearch {
  private vectorizer: TFIDFVectorizer = new TFIDFVectorizer();
  private ruleVectors: RuleVector[] = [];
  private isBuilt = false;

  /**
   * 构建规则向量索引
   */
  buildIndex(rules: Record<string, unknown>[]): void {
    if (rules.length === 0) {
      this.isBuilt = false;
      return;
    }

    // 构建每条规则的文本表示
    const texts = rules.map(rule => this.buildRuleText(rule));

    // 拟合 TF-IDF
    this.vectorizer.fit(texts);

    // 生成向量
    this.ruleVectors = rules.map((doc, idx) => ({
      doc,
      vector: this.vectorizer.transform(texts[idx]),
      text: texts[idx],
    }));

    this.isBuilt = true;
    console.log(`[EmbeddingRetrieval] 向量索引构建完成: ${rules.length} 条规则, 词汇表大小: ${this.vectorizer.getVocabularySize()}`);
  }

  /**
   * 构建规则的文本表示（用于向量化）
   */
  private buildRuleText(rule: Record<string, unknown>): string {
    const parts: string[] = [];

    // 规则名称（权重最高）
    const ruleName = String(rule.ruleName || rule.ruleNameEn || "");
    if (ruleName) parts.push(ruleName, ruleName); // 重复两次增加权重

    // 辖区和资产类别
    const jurisdiction = String(rule.jurisdiction || "");
    const assetClass = String(rule.assetClass || rule.category || "");
    const ruleType = String(rule.ruleType || "");
    if (jurisdiction) parts.push(jurisdiction);
    if (assetClass) parts.push(assetClass);
    if (ruleType) parts.push(ruleType);

    // 内容（主要文本）
    const content = String(rule.content || rule.description || "");
    if (content) parts.push(content.slice(0, 500));

    // 法律依据
    const legalBasis = String(rule.legalBasis || "");
    if (legalBasis) parts.push(legalBasis);

    // 标签
    const tags = Array.isArray(rule.tags) ? rule.tags.join(" ") : "";
    if (tags) parts.push(tags);

    // 所有权要求
    const ownerReqs = rule.ownershipRequirements as Record<string, unknown> | undefined;
    if (ownerReqs) {
      const docs = Array.isArray(ownerReqs.proofDocuments)
        ? ownerReqs.proofDocuments.join(" ")
        : "";
      if (docs) parts.push(docs.slice(0, 200));
    }

    return parts.join(" ");
  }

  /**
   * 语义搜索
   */
  search(query: string, topK = 5, minScore = 0.1): Array<{ doc: Record<string, unknown>; score: number }> {
    if (!this.isBuilt || this.ruleVectors.length === 0) {
      return [];
    }

    const queryVector = this.vectorizer.transform(query);

    // 计算所有规则的相似度
    const scored = this.ruleVectors.map(rv => ({
      doc: rv.doc,
      score: cosineSimilarity(queryVector, rv.vector),
    }));

    // 过滤低分并排序
    return scored
      .filter(item => item.score >= minScore)
      .sort((a, b) => b.score - a.score)
      .slice(0, topK);
  }

  isReady(): boolean {
    return this.isBuilt;
  }
}

// ─── 全局向量检索引擎实例 ─────────────────────────────────────────

const globalVectorEngine = new InMemoryVectorSearch();
let lastIndexBuildTime = 0;
const INDEX_REBUILD_INTERVAL = 5 * 60 * 1000; // 5分钟重建一次

// ─── MongoDB 连接 ─────────────────────────────────────────────────

const MONGO_URL = process.env.NAC_MONGO_URL || "mongodb://root:idP0ZaRGyLsTUA3a@localhost:27017/nac_knowledge_engine?authSource=admin";
const DB_NAME = "nac_knowledge_engine";
const COLLECTION_NAME = "compliance_rules";

async function getCollection(): Promise<Collection<Document>> {
  const client = new MongoClient(MONGO_URL);
  await client.connect();
  return client.db(DB_NAME).collection(COLLECTION_NAME);
}

// ─── 向量索引构建 ─────────────────────────────────────────────────

/**
 * 从 MongoDB 加载所有规则并构建向量索引
 */
export async function buildVectorIndex(): Promise<void> {
  const now = Date.now();
  if (globalVectorEngine.isReady() && now - lastIndexBuildTime < INDEX_REBUILD_INTERVAL) {
    return; // 索引仍然有效
  }

  const client = new MongoClient(MONGO_URL);
  try {
    await client.connect();
    const collection = client.db(DB_NAME).collection(COLLECTION_NAME);

    // 加载所有规则
    const rules = await collection.find({}).toArray();

    if (rules.length === 0) {
      console.log("[EmbeddingRetrieval] 知识库为空，跳过向量索引构建");
      return;
    }

    // 构建向量索引
    globalVectorEngine.buildIndex(rules as unknown as Record<string, unknown>[]);
    lastIndexBuildTime = now;

    console.log(`[EmbeddingRetrieval] 向量索引构建完成，共 ${rules.length} 条规则`);
  } catch (e) {
    console.error(`[EmbeddingRetrieval] 向量索引构建失败: ${(e as Error).message}`);
  } finally {
    await client.close();
  }
}

// ─── 语义检索主函数 ───────────────────────────────────────────────

/**
 * 语义检索：基于 TF-IDF 向量相似度
 *
 * @param query 查询文本
 * @param options 检索选项
 * @returns 语义相关的规则列表
 */
export async function semanticSearch(
  query: string,
  options: {
    topK?: number;
    minScore?: number;
    jurisdiction?: string;
    assetClass?: string;
    ruleType?: string;
  } = {}
): Promise<SemanticSearchResult[]> {
  const { topK = 5, minScore = 0.05, jurisdiction, assetClass, ruleType } = options;

  // 确保向量索引已构建
  await buildVectorIndex();

  if (!globalVectorEngine.isReady()) {
    console.log("[EmbeddingRetrieval] 向量引擎未就绪，返回空结果");
    return [];
  }

  // 构建增强查询（加入辖区和资产类别信息）
  let enhancedQuery = query;
  if (jurisdiction) enhancedQuery += ` ${jurisdiction}`;
  if (assetClass) enhancedQuery += ` ${assetClass}`;
  if (ruleType) enhancedQuery += ` ${ruleType}`;

  // 执行向量搜索
  let results = globalVectorEngine.search(enhancedQuery, topK * 3, minScore);

  // 后过滤（辖区/资产类别/规则类型）
  if (jurisdiction) {
    const jurisdictionResults = results.filter(r => {
      const j = String(r.doc.jurisdiction || "").toUpperCase();
      return j === jurisdiction.toUpperCase() || j === "GLOBAL";
    });
    // 如果辖区过滤后结果太少，保留全局规则
    if (jurisdictionResults.length >= 2) {
      results = jurisdictionResults;
    }
  }

  if (assetClass) {
    const assetResults = results.filter(r => {
      const a = String(r.doc.assetClass || r.doc.category || "").toLowerCase();
      return a.includes(assetClass.toLowerCase()) || a === "all" || !a;
    });
    if (assetResults.length >= 2) {
      results = assetResults;
    }
  }

  if (ruleType) {
    const typeResults = results.filter(r => {
      const t = String(r.doc.ruleType || "").toLowerCase();
      return t.includes(ruleType.toLowerCase());
    });
    if (typeResults.length >= 1) {
      results = typeResults;
    }
  }

  // 取前 topK 条
  results = results.slice(0, topK);

  // 格式化结果
  return results.map(r => {
    const doc = r.doc;
    const rawScore = r.score;

    // 将相似度分数映射到 0.4-1.0 范围（避免低分显示）
    const normalizedScore = 0.4 + rawScore * 0.6;
    const safeScore = isNaN(normalizedScore) ? 0.5 : Math.min(1.0, Math.max(0.0, normalizedScore));

    return {
      ruleId: String(doc.ruleId || doc._id || ""),
      ruleName: String(doc.ruleName || doc.ruleNameEn || "未命名规则"),
      jurisdiction: String(doc.jurisdiction || "未知"),
      assetClass: String(doc.assetClass || doc.category || "通用"),
      ruleType: String(doc.ruleType || "compliance_general"),
      content: String(doc.content || doc.description || "").slice(0, 800),
      legalBasis: doc.legalBasis ? String(doc.legalBasis) : undefined,
      ownershipRequirements: doc.ownershipRequirements as Record<string, unknown> | undefined,
      tradingRequirements: doc.tradingRequirements as Record<string, unknown> | undefined,
      score: safeScore,
      similarityScore: rawScore,
      sourceUrl: doc.sourceUrl ? String(doc.sourceUrl) : undefined,
      tags: Array.isArray(doc.tags) ? doc.tags.map(String) : undefined,
      complianceLevel: doc.complianceLevel ? String(doc.complianceLevel) : undefined,
    };
  });
}

/**
 * 混合检索：结合语义检索和关键词检索，取最优结果
 *
 * @param query 查询文本
 * @param keywordResults 关键词检索结果（来自 ragRetrieval.ts）
 * @param options 检索选项
 * @returns 融合后的检索结果
 */
export async function hybridSearch(
  query: string,
  keywordResults: Array<{ ruleId: string; score: number; [key: string]: unknown }>,
  options: {
    topK?: number;
    jurisdiction?: string;
    assetClass?: string;
    ruleType?: string;
    semanticWeight?: number; // 语义检索权重（0-1），默认 0.6
  } = {}
): Promise<SemanticSearchResult[]> {
  const { topK = 5, semanticWeight = 0.6 } = options;
  const keywordWeight = 1 - semanticWeight;

  // 执行语义检索
  const semanticResults = await semanticSearch(query, {
    topK: topK * 2,
    ...options,
  });

  // 构建关键词结果的 Map（ruleId -> score）
  const keywordScoreMap = new Map<string, number>();
  for (const r of keywordResults) {
    keywordScoreMap.set(r.ruleId, isNaN(r.score) ? 0.5 : r.score);
  }

  // 构建语义结果的 Map（ruleId -> result）
  const semanticMap = new Map<string, SemanticSearchResult>();
  for (const r of semanticResults) {
    semanticMap.set(r.ruleId, r);
  }

  // 融合分数
  const allRuleIdsSet = new Set([
    ...semanticResults.map(r => r.ruleId),
    ...keywordResults.map(r => r.ruleId),
  ]);
  const allRuleIds = Array.from(allRuleIdsSet);

  const fusedResults: Array<SemanticSearchResult & { fusedScore: number }> = [];

  for (const ruleId of allRuleIds) {
    const semanticResult = semanticMap.get(ruleId);
    const semanticScore = semanticResult?.similarityScore || 0;
    const keywordScore = keywordScoreMap.get(ruleId) || 0;

    // 加权融合
    const fusedScore = semanticScore * semanticWeight + keywordScore * keywordWeight;

    if (semanticResult) {
      fusedResults.push({
        ...semanticResult,
        score: 0.4 + fusedScore * 0.6, // 归一化到 0.4-1.0
        fusedScore,
      });
    }
  }

  // 排序并返回
  return fusedResults
    .sort((a, b) => b.fusedScore - a.fusedScore)
    .slice(0, topK)
    .map(({ fusedScore: _fs, ...rest }) => rest);
}

/**
 * 强制重建向量索引（用于知识库更新后）
 */
export async function rebuildVectorIndex(): Promise<{ success: boolean; rulesIndexed: number }> {
  lastIndexBuildTime = 0; // 强制重建
  await buildVectorIndex();
  return {
    success: globalVectorEngine.isReady(),
    rulesIndexed: globalVectorEngine.isReady() ? 1 : 0,
  };
}

/**
 * 获取向量检索引擎状态
 */
export function getEmbeddingStatus(): {
  isReady: boolean;
  lastBuildTime: Date | null;
  engine: string;
} {
  return {
    isReady: globalVectorEngine.isReady(),
    lastBuildTime: lastIndexBuildTime > 0 ? new Date(lastIndexBuildTime) : null,
    engine: "TF-IDF InMemory (v1.0)",
  };
}