// HybridSearchService.ts
import type OpenAI from 'openai';
import type { Db, Document } from 'mongodb';

import { cosine, norm } from '../math/vector';
import type { DbManager } from '../db/db-manager';
import type {
  EmbeddingsStore,
  EmbeddingEntry,
} from '../embeddings/embeddings-store';
import type { ShapesMap } from '../shapes/shapes';

export type TextCandidate = {
  collection: string;
  _id_raw: any;
  _id_str: string;
};

export type ScoredCandidate = TextCandidate & {
  score: number;
};

export type HybridSearchDeps = {
  dbManager: DbManager;
  openai: OpenAI;
  embedModel: string;
  embeddingsStore: EmbeddingsStore;
  shapes: ShapesMap;
  maxTextHits?: number;
};

export class HybridSearchService {
  private dbManager: DbManager;
  private openai: OpenAI;
  private embedModel: string;
  private embeddingsStore: EmbeddingsStore;
  private shapes: ShapesMap;
  private maxTextHits: number;

  constructor({
    dbManager,
    openai,
    embedModel,
    embeddingsStore,
    shapes,
    maxTextHits = 200,
  }: HybridSearchDeps) {
    this.dbManager = dbManager;
    this.openai = openai;
    this.embedModel = embedModel;
    this.embeddingsStore = embeddingsStore;
    this.shapes = shapes;
    this.maxTextHits = maxTextHits;
  }

  fetchTextCandidatesForCollection = async (
    db: Db,
    name: string,
    query: string,
  ): Promise<TextCandidate[]> => {
    // Keep schema loose to avoid MongoDB driver _id inference issues (mongodb@7 + TS4.9)
    const coll = db.collection<Document>(name);

    try {
      const hits = await coll
        .find({ $text: { $search: query } } as any, {
          projection: { score: { $meta: 'textScore' }, _id: 1 } as any,
          sort: { score: { $meta: 'textScore' } } as any,
          limit: this.maxTextHits,
        })
        .toArray();

      return hits.map((h: any) => ({
        collection: name,
        _id_raw: h._id,
        _id_str: String(h._id),
      }));
    } catch {
      return [];
    }
  };

  hybridSearch = async (query: string, kPerCollection = 6): Promise<any[]> => {
    const db = await this.dbManager.getDb();
    const ACTIVE = this.embeddingsStore.getActive();

    // 1) parallel text candidates
    const lists = await Promise.all(
      ACTIVE.map((name) =>
        this.fetchTextCandidatesForCollection(db, name, query),
      ),
    );

    // flatten + dedupe
    const candidates: TextCandidate[] = [];
    const seen = new Set<string>();

    for (const list of lists) {
      for (const c of list) {
        const key = `${c.collection}:${c._id_str}`;
        if (seen.has(key)) continue;
        seen.add(key);
        candidates.push(c);
      }
    }

    if (!candidates.length) return [];

    // 2) embed query once
    const er = await this.openai.embeddings.create({
      model: this.embedModel,
      input: query,
    });
    const qv = er.data[0].embedding;
    const qn = norm(qv);

    // 3) semantic rerank
    const scored: ScoredCandidate[] = [];
    for (const c of candidates) {
      const ent: EmbeddingEntry | undefined = this.embeddingsStore
        .getMap(c.collection)
        .get(c._id_str);
      if (!ent) continue;
      scored.push({ ...c, score: cosine(qv, ent.v, qn, ent.n) });
    }

    if (!scored.length) return [];

    // 4) top K per collection
    const byCol = new Map<string, ScoredCandidate[]>();
    for (const s of scored) {
      if (!byCol.has(s.collection)) byCol.set(s.collection, []);
      byCol.get(s.collection)!.push(s);
    }
    for (const arr of byCol.values()) arr.sort((a, b) => b.score - a.score);

    const topPer: ScoredCandidate[] = [];
    for (const [, arr] of byCol) topPer.push(...arr.slice(0, kPerCollection));

    // 5) hydrate docs
    const groupedIds = new Map<string, any[]>();
    for (const t of topPer) {
      if (!groupedIds.has(t.collection)) groupedIds.set(t.collection, []);
      groupedIds.get(t.collection)!.push(t._id_raw);
    }

    const results: any[] = [];
    for (const [name, ids] of groupedIds) {
      const shape = (this.shapes as any)[name] as
        | ShapesMap[keyof ShapesMap]
        | undefined;
      const projection = shape?.projection ?? {};

      // ✅ Key fix:
      // mongodb@7 types sometimes infer `_id` operator value type as `never` when collection schema is generic.
      // Casting the filter to `any` keeps runtime correct and fixes TS2769.
      const docs = await db
        .collection<Document>(name)
        .find({ _id: { $in: ids } } as any, { projection: projection as any })
        .toArray();

      for (const d of docs) {
        const h = shape?.hydrate
          ? shape.hydrate(d as any)
          : { collection: name, _id: String((d as any)._id), ...(d as any) };
        results.push(h);
      }
    }

    // preserve semantic order
    const order = new Map<string, number>(
      topPer.map((t, i) => [`${t.collection}:${String(t._id_raw)}`, i]),
    );
    results.sort(
      (a, b) =>
        (order.get(`${a.collection}:${a._id}`) ?? 0) -
        (order.get(`${b.collection}:${b._id}`) ?? 0),
    );

    return results;
  };
}