import indexData from './index.json';

export type DocChunk = {
  id: string;
  table: string;
  rowId: string;
  chunkIndex: number;
  text: string;
  meta: Record<string, any>;
  vector: number[];
};

export type AskInput = {
  q: string;
  tables?: string[]; // optional table filter
  topK?: number; // default from env
  dedupeByRow?: boolean; // default true
};

export type AskOutput = {
  answer: string;
  sources: Array<{
    id: string;
    table: string;
    rowId: string;
    score: number;
    title?: string;
    thumbnail?: string;
    venue?: string;
    address?: string;
    slug?: string;
  }>;
};

type ServiceOpts = {
  // indexPath is removed since we are importing the JSON directly
  ollamaUrl: string;
  embedModel: string;
  chatModel: string;
  topK: number;
  requestTimeoutMs: number;
};

export class AskChatService {
  private index: DocChunk[];
  private opts: ServiceOpts;

  constructor(opts: ServiceOpts) {
    this.opts = opts;

    // Assign the imported JSON directly to our index array
    this.index = indexData as DocChunk[];
    console.log(
      `[AskChatService] Loaded chunks=${this.index.length} via static import`,
    );
  }

  async ask(input: AskInput): Promise<AskOutput> {
    const q = input.q?.trim();
    if (!q) throw new Error('Missing query');

    const topK = Number.isFinite(input.topK)
      ? Math.max(1, Number(input.topK))
      : this.opts.topK;
    const dedupeByRow = input.dedupeByRow ?? true;

    const qVec = await this.embedQuery(q);

    const hits = this.searchTopK(qVec, topK, {
      tables: input.tables,
      dedupeByRow,
    });

    const context = this.buildContext(hits);

    const answer = await this.chat(q, context);

    return {
      answer,
      sources: hits.map((h) => ({
        id: h.chunk.id,
        table: h.chunk.table,
        rowId: h.chunk.rowId,
        score: h.score,
        title: h.chunk.meta?.title,
        thumbnail: h.chunk.meta?.thumbnail,
        venue: h.chunk.meta?.venue,
        address: h.chunk.meta?.formatted_address,
        slug: h.chunk.meta?.slug,
      })),
    };
  }

  async askAuto(q: string) {
    // Step 1: search broadly (all tables), slightly higher K to classify
    const initialTopK = Number(process.env.AUTO_TOP_K ?? 10);

    // Use your existing ask() but without tables filter
    // dedupeByRow true keeps context clean
    const broad = await this.ask({
      q,
      topK: initialTopK,
      tables: undefined,
      dedupeByRow: true,
    });

    // Step 2 (optional): if you want “more accurate data accordingly”
    // pick dominant table among top hits, then re-run with that table only
    const dominantTable = this.mostCommon(broad.sources.map((s) => s.table));

    if (dominantTable) {
      const refinedTopK = Number(process.env.REFINE_TOP_K ?? 6);
      const refined = await this.ask({
        q,
        topK: refinedTopK,
        tables: [dominantTable],
        dedupeByRow: true,
      });

      return {
        ...refined,
        inferredTable: dominantTable,
        mode: 'refined',
      };
    }

    return { ...broad, inferredTable: null, mode: 'broad' };
  }

  async askAutoFast(q: string) {
    const topK = Number(process.env.TOP_K ?? 6);
    const probeK = Number(process.env.PROBE_K ?? 12); // slightly bigger to infer table

    // 1) embed once
    const qVec = await this.embedQuery(q);

    // 2) search once (no table filter)
    const broadHits = this.searchTopKFromVector(qVec, probeK, {
      tables: undefined,
      dedupeByRow: true,
    });

    // 3) infer dominant table from top hits
    const dominantTable = this.fastmostCommon(
      broadHits.map((h) => h.chunk.table),
    );

    // 4) refine from the SAME hits (no second embed/search)
    const refinedHits = dominantTable
      ? broadHits.filter((h) => h.chunk.table === dominantTable).slice(0, topK)
      : broadHits.slice(0, topK);

    const context = this.buildContext(refinedHits);

    // 5) chat once
    const answer = await this.chat(q, context);

    return {
      answer,
      inferredTable: dominantTable,
      sources: refinedHits.map((h) => ({
        table: h.chunk.table,
        rowId: h.chunk.rowId,
        score: h.score,
        title: h.chunk.meta?.title,
        thumbnail: h.chunk.meta?.thumbnail,
        venue: h.chunk.meta?.venue,
        address: h.chunk.meta?.formatted_address,
        slug: h.chunk.meta?.slug,
      })),
    };
  }

  fastmostCommon(items: (string | undefined)[]) {
    const m = new Map<string, number>();
    for (const it of items) if (it) m.set(it, (m.get(it) ?? 0) + 1);
    let best: string | null = null;
    let bestN = 0;
    for (const [k, n] of m)
      if (n > bestN) {
        best = k;
        bestN = n;
      }
    return best;
  }

  mostCommon(items: Array<string | undefined | null>): string | null {
    const freq = new Map<string, number>();
    for (const it of items) {
      if (!it) continue;
      freq.set(it, (freq.get(it) ?? 0) + 1);
    }
    let best: string | null = null;
    let bestN = 0;
    for (const [k, n] of freq) {
      if (n > bestN) {
        best = k;
        bestN = n;
      }
    }
    return best;
  }

  // ---------- core helpers ----------

  private dot(a: number[], b: number[]) {
    let s = 0;
    const n = Math.min(a.length, b.length);
    for (let i = 0; i < n; i++) s += a[i] * b[i];
    return s;
  }

  private searchTopK(
    qVec: number[],
    topK: number,
    opts: { tables?: string[]; dedupeByRow: boolean },
  ): Array<{ score: number; chunk: DocChunk }> {
    const tableSet = opts.tables?.length ? new Set(opts.tables) : null;

    const scored: Array<{ score: number; chunk: DocChunk }> = [];
    for (const ch of this.index) {
      if (tableSet && !tableSet.has(ch.table)) continue;
      scored.push({ score: this.dot(qVec, ch.vector), chunk: ch });
    }

    scored.sort((a, b) => b.score - a.score);

    if (!opts.dedupeByRow) return scored.slice(0, topK);

    // keep best chunk per (table,rowId) to reduce repeated context
    const seen = new Set<string>();
    const out: Array<{ score: number; chunk: DocChunk }> = [];
    for (const s of scored) {
      const key = `${s.chunk.table}:${s.chunk.rowId}`;
      if (seen.has(key)) continue;
      seen.add(key);
      out.push(s);
      if (out.length >= topK) break;
    }
    return out;
  }

  private searchTopKFromVector(
    qVec: number[],
    topK: number,
    opts: { tables?: string[]; dedupeByRow: boolean },
  ): Array<{ score: number; chunk: DocChunk }> {
    const tableSet = opts.tables?.length ? new Set(opts.tables) : null;
    const K = Math.max(0, topK | 0);
    if (K === 0) return [];

    type Hit = { score: number; chunk: DocChunk };

    // ---- tiny min-heap helpers (keeps smallest score at heap[0]) ----
    const heap: Hit[] = [];

    const siftUp = (i: number) => {
      while (i > 0) {
        const p = (i - 1) >> 1;
        if (heap[p].score <= heap[i].score) break;
        [heap[p], heap[i]] = [heap[i], heap[p]];
        i = p;
      }
    };

    const siftDown = (i: number) => {
      for (;;) {
        const l = i * 2 + 1;
        const r = l + 1;
        let smallest = i;

        if (l < heap.length && heap[l].score < heap[smallest].score)
          smallest = l;
        if (r < heap.length && heap[r].score < heap[smallest].score)
          smallest = r;

        if (smallest === i) break;
        [heap[i], heap[smallest]] = [heap[smallest], heap[i]];
        i = smallest;
      }
    };

    const heapPush = (h: Hit) => {
      heap.push(h);
      siftUp(heap.length - 1);
    };

    const heapReplaceRoot = (h: Hit) => {
      heap[0] = h;
      siftDown(0);
    };

    const consider = (h: Hit) => {
      if (heap.length < K) {
        heapPush(h);
        return;
      }
      // if better than current worst(best-of-topK min), replace
      if (h.score > heap[0].score) heapReplaceRoot(h);
    };

    // ---- main scan ----
    if (!opts.dedupeByRow) {
      for (const ch of this.index) {
        if (tableSet && !tableSet.has(ch.table)) continue;
        const score = this.dot(qVec, ch.vector);
        consider({ score, chunk: ch });
      }

      // heap has up to K best; return sorted desc
      heap.sort((a, b) => b.score - a.score);
      return heap;
    }

    // dedupe mode: keep best per (table,rowId)
    const bestByRow = new Map<string, Hit>();

    for (const ch of this.index) {
      if (tableSet && !tableSet.has(ch.table)) continue;

      const score = this.dot(qVec, ch.vector);
      const key = `${ch.table}:${ch.rowId}`;
      const prev = bestByRow.get(key);

      if (!prev || score > prev.score) {
        bestByRow.set(key, { score, chunk: ch });
      }
    }

    // select topK from the deduped set
    for (const hit of bestByRow.values()) {
      consider(hit);
    }

    heap.sort((a, b) => b.score - a.score);
    return heap;
  }

  private buildContext(hits: Array<{ score: number; chunk: DocChunk }>) {
    // Keep it compact for a 1B model
    return hits
      .map((h, i) => {
        const m = h.chunk.meta ?? {};
        const lines = [
          `### RESULT ${i + 1}`,
          `TABLE: ${h.chunk.table}`,
          `ROW_ID: ${h.chunk.rowId}`,
          m.title ? `TITLE: ${m.title}` : '',
          m.slug ? `SLUG: ${m.slug}` : '',
          m.venue ? `VENUE: ${m.venue}` : '',
          m.formatted_address ? `ADDRESS: ${m.formatted_address}` : '',
          m.open_days || m.opened_at || m.closed_at
            ? `HOURS: ${[
                m.open_days,
                m.opened_at && m.closed_at
                  ? `${m.opened_at}-${m.closed_at}`
                  : '',
              ]
                .filter(Boolean)
                .join(' ')}`
            : '',
          m.fee ? `FEE: ${m.fee}` : '',
          m.thumbnail ? `THUMBNAIL: ${m.thumbnail}` : '',
          `TEXT:\n${h.chunk.text}`,
        ].filter(Boolean);
        return lines.join('\n');
      })
      .join('\n\n');
  }

  private async embedQuery(q: string): Promise<number[]> {
    const res = await this.fetchWithTimeout(
      `${this.opts.ollamaUrl}/api/embed`,
      {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          model: this.opts.embedModel,
          input: q,
          truncate: true,
        }),
      },
    );

    if (!res.ok) {
      throw new Error(
        `Ollama embed failed (${res.status}): ${await res.text()}`,
      );
    }
    const json = await res.json();
    const vec = (json.embeddings as number[][])?.[0];
    if (!vec?.length) throw new Error('No embedding returned from Ollama');
    return vec;
  }

  private async chat(question: string, context: string): Promise<string> {
    const res = await this.fetchWithTimeout(`${this.opts.ollamaUrl}/api/chat`, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
        model: this.opts.chatModel,
        stream: false,
        messages: [
          {
            role: 'system',
            content:
              'You are DoseeBot, a travel & culture assistant. ' +
              'Use ONLY the provided CONTEXT. If the answer is not in context, say you do not know and suggest what to add.',
          },
          {
            role: 'user',
            content: `QUESTION:\n${question}\n\nCONTEXT:\n${context}\n\nReturn a helpful answer. If relevant, include venue, hours, fee, and share any thumbnail URL present.`,
          },
        ],
      }),
    });

    if (!res.ok) {
      console.log(await res.json());
      throw new Error(
        `Ollama chat failed (${res.status}): ${await res.text()}`,
      );
    }
    const json = await res.json();
    return json?.message?.content ?? '';
  }

  private async fetchWithTimeout(url: string, init: RequestInit) {
    const controller = new AbortController();
    const timer = setTimeout(
      () => controller.abort(),
      this.opts.requestTimeoutMs,
    );
    try {
      return await fetch(url, { ...init, signal: controller.signal });
    } finally {
      clearTimeout(timer);
    }
  }
}