import { tool } from 'ai'
import { z } from 'zod'
import { toolRegistry } from './registry'
import { embeddingService } from '@/lib/ai/services/embedding.service'
import prisma from '@/lib/prisma'

toolRegistry.register({
  name: 'document_search',
  description: 'Search within PDF documents attached to notes. Returns relevant passages with page numbers and source document info.',
  isInternal: true,
  buildTool: (ctx) =>
    tool({
      description: `Search within PDF documents attached to the user's notes.
Returns matching passages with page numbers, chunk content, and the source note/document info.
Use this when the user asks about specific documents, PDFs, or attached files.`,
      inputSchema: z.object({
        query: z.string().describe('The search query to find relevant passages in documents'),
        noteId: z.string().optional().describe('Optional: restrict search to attachments of a specific note'),
        limit: z.number().optional().describe('Max results to return (default 5)').default(5),
      }),
      execute: async ({ query, noteId, limit = 5 }) => {
        try {
          const queryEmbedding = await embeddingService.generateEmbedding(query)
          const vectorStr = embeddingService.toVectorString(queryEmbedding.embedding)

          let noteFilter = ''
          const params: any[] = [vectorStr, limit, ctx.userId]

          if (noteId) {
            noteFilter = `AND na."noteId" = $4`
            params.push(noteId)
          }

          const results = await prisma.$queryRawUnsafe(
            `SELECT
              dc.id as "chunkId",
              dc.content,
              dc."pageNumber",
              dc."chunkIndex",
              na.id as "attachmentId",
              na."fileName",
              na."pageCount",
              na."noteId",
              n.title as "noteTitle"
            FROM "DocumentChunk" dc
            JOIN "NoteAttachment" na ON na.id = dc."attachmentId"
            JOIN "Note" n ON n.id = na."noteId"
            WHERE dc."embedding" IS NOT NULL
              AND na.status = 'ready'
              AND n."trashedAt" IS NULL
              AND n."userId" = $3
              ${noteFilter}
            ORDER BY dc."embedding" <=> $1::vector
            LIMIT $2`,
            ...params
          ) as any[]

          if (!results.length) return { results: [], message: 'No matching documents found' }

          return results.map(r => ({
            content: r.content.substring(0, 600),
            pageNumber: r.pageNumber,
            chunkIndex: r.chunkIndex,
            fileName: r.fileName,
            noteId: r.noteId,
            noteTitle: r.noteTitle || 'Untitled',
          }))
        } catch (e: any) {
          return { error: `Document search failed: ${e.message}` }
        }
      },
    }),
})