import { randomUUID } from 'node:crypto';
import type { IQueryEmbedding } from '../interfaces/query-embedding.js';
import type { IRag, IRagBackendWriter } from '../interfaces/rag.js';
import type {
  CallOptions,
  RagMetadata,
  RagResult,
  Result,
} from '../interfaces/types.js';
import { RagError } from '../interfaces/types.js';
import type { IDocumentEnricher, IQueryPreprocessor } from './preprocessor.js';

// ---------------------------------------------------------------------------
// Private helpers
// ---------------------------------------------------------------------------

function tokenize(text: string): string[] {
  return text
    .toLowerCase()
    .replace(/[^a-z0-9\s]/g, ' ')
    .split(/\s+/)
    .filter((t) => t.length > 1);
}

function embed(text: string): Map<string, number> {
  const tokens = tokenize(text);
  const freq = new Map<string, number>();
  for (const t of tokens) freq.set(t, (freq.get(t) ?? 0) + 1);
  const norm = Math.sqrt([...freq.values()].reduce((s, v) => s + v * v, 0));
  if (norm === 0) return freq;
  for (const [k, v] of freq) freq.set(k, v / norm);
  return freq;
}

function cosineSimilarity(
  a: Map<string, number>,
  b: Map<string, number>,
): number {
  let dot = 0;
  for (const [term, wa] of a) {
    const wb = b.get(term);
    if (wb !== undefined) dot += wa * wb;
  }
  return dot; // both are unit vectors → dot = cosine
}

// ---------------------------------------------------------------------------
// Config
// ---------------------------------------------------------------------------

export interface InMemoryRagConfig {
  /** Cosine similarity above which upsert updates existing record. Default: 0.92 */
  dedupThreshold?: number;
  /** Namespace for this store. Records with different namespace are invisible to query. */
  namespace?: string;
  /** Query preprocessors (translate, expand, etc.). Applied in order before embedding. */
  queryPreprocessors?: IQueryPreprocessor[];
  /** Document enrichers. Applied in order before embedding on upsert. */
  documentEnrichers?: IDocumentEnricher[];
}

// ---------------------------------------------------------------------------
// Internal record
// ---------------------------------------------------------------------------

interface StoredRecord {
  id: string;
  text: string;
  embedding: Map<string, number>;
  metadata: RagMetadata;
}

// ---------------------------------------------------------------------------
// InMemoryRag
// ---------------------------------------------------------------------------

export class InMemoryRag implements IRag {
  private records: StoredRecord[] = [];
  private readonly dedupThreshold: number;
  private readonly namespace?: string;
  private readonly queryPreprocessors: IQueryPreprocessor[];
  private readonly documentEnrichers: IDocumentEnricher[];

  constructor(config?: InMemoryRagConfig) {
    this.dedupThreshold = config?.dedupThreshold ?? 0.92;
    this.namespace = config?.namespace;
    this.queryPreprocessors = config?.queryPreprocessors ?? [];
    this.documentEnrichers = config?.documentEnrichers ?? [];
  }

  async upsert(
    text: string,
    metadata: RagMetadata,
    options?: CallOptions,
  ): Promise<Result<void, RagError>> {
    if (options?.signal?.aborted) {
      return { ok: false, error: new RagError('Aborted', 'ABORTED') };
    }

    let enrichedText = text;
    for (const enricher of this.documentEnrichers) {
      const eResult = await enricher.enrich(enrichedText, options);
      if (eResult.ok) enrichedText = eResult.value;
    }

    const embedding = embed(enrichedText);
    const effectiveNamespace = metadata.namespace ?? this.namespace;
    const resolvedMetadata: RagMetadata = {
      ...metadata,
      namespace: effectiveNamespace,
    };

    // Idempotent upsert: if metadata.id matches, replace in-place
    if (metadata.id) {
      const idx = this.records.findIndex((r) => r.metadata.id === metadata.id);
      if (idx !== -1) {
        this.records[idx].text = enrichedText;
        this.records[idx].embedding = embedding;
        this.records[idx].metadata = {
          ...this.records[idx].metadata,
          ...resolvedMetadata,
        };
        return { ok: true, value: undefined };
      }
    }

    // Filter existing records by same namespace
    const candidates =
      this.namespace !== undefined
        ? this.records.filter((r) => r.metadata.namespace === this.namespace)
        : this.records;

    // Find record with cosine similarity >= dedupThreshold
    let dupRecord: StoredRecord | undefined;
    for (const r of candidates) {
      if (cosineSimilarity(embedding, r.embedding) >= this.dedupThreshold) {
        dupRecord = r;
        break;
      }
    }

    if (dupRecord !== undefined) {
      // Update existing record
      dupRecord.text = enrichedText;
      dupRecord.embedding = embedding;
      dupRecord.metadata = { ...dupRecord.metadata, ...resolvedMetadata };
    } else {
      // Push new record
      this.records.push({
        id: randomUUID(),
        text: enrichedText,
        embedding,
        metadata: resolvedMetadata,
      });
    }

    return { ok: true, value: undefined };
  }

  async query(
    embedding: IQueryEmbedding,
    k: number,
    options?: CallOptions,
  ): Promise<Result<RagResult[], RagError>> {
    if (options?.signal?.aborted) {
      return { ok: false, error: new RagError('Aborted', 'ABORTED') };
    }

    const text = embedding.text;
    let searchText = text;
    for (const pp of this.queryPreprocessors) {
      const ppResult = await pp.process(searchText, options);
      if (ppResult.ok) searchText = ppResult.value;
    }
    const queryEmbedding = embed(searchText);
    const nowSecs = Date.now() / 1000;
    const targetSessionId = options?.ragFilter?.sessionId;

    // Filter: namespace match + TTL not expired + sessionId match
    const candidates = this.records.filter((r) => {
      if (
        this.namespace !== undefined &&
        r.metadata.namespace !== this.namespace
      )
        return false;
      if (r.metadata.ttl !== undefined && r.metadata.ttl < nowSecs)
        return false;
      if (
        targetSessionId !== undefined &&
        r.metadata.sessionId !== targetSessionId
      )
        return false;
      return true;
    });

    // Compute cosine similarity for each candidate
    const scored = candidates.map((r) => ({
      text: r.text,
      metadata: r.metadata,
      score: cosineSimilarity(queryEmbedding, r.embedding),
    }));

    // Sort desc by score, take top k
    scored.sort((a, b) => b.score - a.score);
    const results: RagResult[] = scored.slice(0, k);

    return { ok: true, value: results };
  }

  async getById(
    id: string,
    _options?: CallOptions,
  ): Promise<Result<RagResult | null, RagError>> {
    const record = this.records.find((r) => r.metadata.id === id);
    if (!record) return { ok: true, value: null };
    return {
      ok: true,
      value: { text: record.text, metadata: record.metadata, score: 1 },
    };
  }

  async healthCheck(): Promise<Result<void, RagError>> {
    return { ok: true, value: undefined };
  }

  clear(): void {
    this.records.length = 0;
  }

  writer(): IRagBackendWriter {
    return {
      upsertRaw: async (id, text, metadata, options) => {
        const res = await this.upsert(text, { ...metadata, id }, options);
        return res.ok ? { ok: true, value: undefined } : res;
      },
      deleteByIdRaw: async (id) => {
        const idx = this.records.findIndex((r) => r.metadata.id === id);
        if (idx === -1) return { ok: true, value: false };
        this.records.splice(idx, 1);
        return { ok: true, value: true };
      },
      clearAll: async () => {
        this.records.length = 0;
        return { ok: true, value: undefined };
      },
    };
  }
}
