import "server-only";
import { sql, vec } from "../db";
import { openaiEmbed } from "../ai/openai";
import { recordUsage } from "../usage";
import { chunkText } from "../storage";
import type { ToolContext } from "./index";

/** Index an attachment's extracted text into document_chunks (embeddings + FTS). */
export async function indexAttachment(attachmentId: string) {
  const a = await sql<{ extracted_text: string | null }[]>`SELECT extracted_text FROM attachments WHERE id = ${attachmentId}`;
  const text = a[0]?.extracted_text;
  if (!text || text.length < 40) return 0;
  await sql`DELETE FROM document_chunks WHERE attachment_id = ${attachmentId}`;
  const chunks = chunkText(text).slice(0, 400); // cap ~500k chars per document
  for (let i = 0; i < chunks.length; i += 64) {
    const batch = chunks.slice(i, i + 64);
    let vectors: number[][] | null = null;
    try { const r = await openaiEmbed(batch); vectors = r.vectors; await recordUsage({ provider: "openai", modelId: "text-embedding-3-small", kind: "embedding", usage: { input: r.tokens, output: 0 }, costEstimated: r.tokens * 0.02 / 1e6 }); } catch { vectors = null; }
    for (let j = 0; j < batch.length; j++) {
      const v = vectors?.[j];
      if (v) await sql`INSERT INTO document_chunks (attachment_id, idx, content, embedding) VALUES (${attachmentId}, ${i + j}, ${batch[j]}, ${vec(v)}::vector)`;
      else await sql`INSERT INTO document_chunks (attachment_id, idx, content) VALUES (${attachmentId}, ${i + j}, ${batch[j]})`;
    }
  }
  return chunks.length;
}

export async function searchDocuments(query: string, ctx: Pick<ToolContext, "conversationId" | "projectId">, limit = 6) {
  let qv: number[] | null = null;
  try { qv = (await openaiEmbed([query])).vectors[0]; } catch { qv = null; }
  const rows = await sql<{ id: number; content: string; filename: string; page: number | null; attachment_id: string; score: number }[]>`
    SELECT c.id, c.content, c.page, c.attachment_id, a.filename,
      (${qv ? sql`1 - (c.embedding <=> ${vec(qv)}::vector)` : sql`0`}) * 0.7 + ts_rank(c.tsv, plainto_tsquery('simple', ${query})) * 0.3 AS score
    FROM document_chunks c JOIN attachments a ON a.id = c.attachment_id
    WHERE (a.conversation_id = ${ctx.conversationId} OR (${ctx.projectId ?? null}::text IS NOT NULL AND a.project_id = ${ctx.projectId ?? null}))
    ORDER BY score DESC LIMIT ${limit}`;
  return { query, passages: rows.map(r => ({ file: r.filename, attachmentId: r.attachment_id, page: r.page, score: Number(r.score.toFixed(3)), text: r.content })) };
}