Keep/keep-notes/app/api/chat/route.ts

import { streamText, UIMessage } from 'ai'
import { getChatProvider } from '@/lib/ai/factory'
import { getSystemConfig } from '@/lib/config'
import { semanticSearchService } from '@/lib/ai/services/semantic-search.service'
import { prisma } from '@/lib/prisma'
import { auth } from '@/auth'
import { loadTranslations, getTranslationValue, SupportedLanguage } from '@/lib/i18n'

export const maxDuration = 60

/**
 * Extract text content from a UIMessage's parts array.
 */
function extractTextFromUIMessage(msg: { parts?: Array<{ type: string; text?: string }>; content?: string }): string {
  if (typeof msg.content === 'string') return msg.content
  if (msg.parts && Array.isArray(msg.parts)) {
    return msg.parts
      .filter((p) => p.type === 'text' && typeof p.text === 'string')
      .map((p) => p.text!)
      .join('')
  }
  return ''
}

/**
 * Convert an array of UIMessages (from the client) to CoreMessage[] for streamText.
 */
function toCoreMessages(uiMessages: UIMessage[]): Array<{ role: 'user' | 'assistant'; content: string }> {
  return uiMessages
    .filter((m) => m.role === 'user' || m.role === 'assistant')
    .map((m) => ({
      role: m.role as 'user' | 'assistant',
      content: extractTextFromUIMessage(m),
    }))
    .filter((m) => m.content.length > 0)
}

export async function POST(req: Request) {
  // 1. Auth check
  const session = await auth()
  if (!session?.user?.id) {
    return new Response('Unauthorized', { status: 401 })
  }
  const userId = session.user.id

  // 2. Parse request body — messages arrive as UIMessage[] from DefaultChatTransport
  const body = await req.json()
  const { messages: rawMessages, conversationId, notebookId, language } = body as {
    messages: UIMessage[]
    conversationId?: string
    notebookId?: string
    language?: string
  }

  // Convert UIMessages to CoreMessages for streamText
  const incomingMessages = toCoreMessages(rawMessages)

  // 3. Manage conversation (create or fetch)
  let conversation: { id: string; messages: Array<{ role: string; content: string }> }

  if (conversationId) {
    const existing = await prisma.conversation.findUnique({
      where: { id: conversationId, userId },
      include: { messages: { orderBy: { createdAt: 'asc' } } },
    })
    if (!existing) {
      return new Response('Conversation not found', { status: 404 })
    }
    conversation = existing
  } else {
    const userMessage = incomingMessages[incomingMessages.length - 1]?.content || 'New conversation'
    const created = await prisma.conversation.create({
      data: {
        userId,
        notebookId: notebookId || null,
        title: userMessage.substring(0, 50) + (userMessage.length > 50 ? '...' : ''),
      },
      include: { messages: true },
    })
    conversation = created
  }

  // 4. RAG retrieval
  const currentMessage = incomingMessages[incomingMessages.length - 1]?.content || ''

  // Load translations for the requested language
  const lang = (language || 'en') as SupportedLanguage
  const translations = await loadTranslations(lang)
  const untitledText = getTranslationValue(translations, 'notes.untitled') || 'Untitled'

  // If a notebook is selected, fetch its recent notes directly as context
  // This ensures the AI always has access to the notebook content,
  // even for vague queries like "what's in this notebook?"
  let notebookContext = ''
  if (notebookId) {
    const notebookNotes = await prisma.note.findMany({
      where: {
        notebookId,
        userId,
        trashedAt: null,
      },
      orderBy: { updatedAt: 'desc' },
      take: 20,
      select: { id: true, title: true, content: true, updatedAt: true },
    })
    if (notebookNotes.length > 0) {
      notebookContext = notebookNotes
        .map(n => `NOTE [${n.title || untitledText}] (updated ${n.updatedAt.toLocaleDateString()}):\n${(n.content || '').substring(0, 1500)}`)
        .join('\n\n---\n\n')
    }
  }

  // Also run semantic search for the specific query
  const searchResults = await semanticSearchService.search(currentMessage, {
    notebookId,
    limit: notebookId ? 10 : 5,
    threshold: notebookId ? 0.3 : 0.5,
    defaultTitle: untitledText,
  })

  const searchNotes = searchResults
    .map((r) => `NOTE [${r.title || untitledText}]: ${r.content}`)
    .join('\n\n---\n\n')

  // Combine: full notebook context + semantic search results (deduplicated)
  const contextNotes = [notebookContext, searchNotes].filter(Boolean).join('\n\n---\n\n')

  // 5. System prompt synthesis with RAG context
  // Language-aware prompts to avoid forcing French responses
  // Note: lang is already declared above when loading translations
  const promptLang: Record<string, { contextWithNotes: string; contextNoNotes: string; system: string }> = {
    en: {
      contextWithNotes: `## User's notes\n\n${contextNotes}\n\nWhen using info from the notes above, cite the source note title in parentheses, e.g.: "Deployment is done via Docker (💻 Development Guide)". Don't copy word for word — rephrase. If the notes don't cover the topic, say so and supplement with your general knowledge.`,
      contextNoNotes: "No relevant notes found for this question. Answer with your general knowledge.",
      system: `You are the AI assistant of Memento. The user asks you questions about their projects, technical docs, and notes. You must respond in a structured and helpful way.

## Format rules
- Use markdown freely: headings (##, ###), lists, code blocks, bold, tables — anything that makes the response readable.
- Structure your response with sections for technical questions or complex topics.
- For simple, short questions, a direct paragraph is enough.

## Tone rules
- Natural tone, neither corporate nor too casual.
- No unnecessary intro phrases ("Here's what I found", "Based on your notes"). Answer directly.
- No upsell questions at the end ("Would you like me to...", "Do you want..."). If you have useful additional info, just give it.
- If the user says "Momento" they mean Memento (this app).`,
    },
    fr: {
      contextWithNotes: `## Notes de l'utilisateur\n\n${contextNotes}\n\nQuand tu utilises une info venant des notes ci-dessus, cite le titre de la note source entre parenthèses, ex: "Le déploiement se fait via Docker (💻 Development Guide)". Ne recopie pas mot pour mot — reformule. Si les notes ne couvrent pas le sujet, dis-le et complète avec tes connaissances générales.`,
      contextNoNotes: "Aucune note pertinente trouvée pour cette question. Réponds avec tes connaissances générales.",
      system: `Tu es l'assistant IA de Memento. L'utilisateur te pose des questions sur ses projets, sa doc technique, ses notes. Tu dois répondre de façon structurée et utile.

## Règles de format
- Utilise le markdown librement : titres (##, ###), listes, code blocks, gras, tables — tout ce qui rend la réponse lisible.
- Structure ta réponse avec des sections quand c'est une question technique ou un sujet complexe.
- Pour les questions simples et courtes, un paragraphe direct suffit.

## Règles de ton
- Ton naturel, ni corporate ni trop familier.
- Pas de phrase d'intro inutile ("Voici ce que j'ai trouvé", "Basé sur vos notes"). Réponds directement.
- Pas de question upsell à la fin ("Souhaitez-vous que je...", "Acceptez-vous que..."). Si tu as une info complémentaire utile, donne-la.
- Si l'utilisateur dit "Momento" il parle de Memento (cette application).`,
    },
    fa: {
      contextWithNotes: `## یادداشت‌های کاربر\n\n${contextNotes}\n\nهنگام استفاده از اطلاعات یادداشت‌های بالا، عنوان یادداشت منبع را در پرانتز ذکر کنید. کپی نکنید — بازنویسی کنید. اگر یادداشت‌ها موضوع را پوشش نمی‌دهند، بگویید و با دانش عمومی خود تکمیل کنید.`,
      contextNoNotes: "هیچ یادداشت مرتبطی برای این سؤال یافت نشد. با دانش عمومی خود پاسخ دهید.",
      system: `شما دستیار هوش مصنوعی Memento هستید. کاربر از شما درباره پروژه‌ها، مستندات فنی و یادداشت‌هایش سؤال می‌کند. باید به شکلی ساختاریافته و مفید پاسخ دهید.

## قوانین قالب‌بندی
- از مارک‌داون آزادانه استفاده کنید: عناوین (##, ###)، لیست‌ها، بلوک‌های کد، پررنگ، جداول.
- برای سؤالات فنی یا موضوعات پیچیده، پاسخ خود را بخش‌بندی کنید.
- برای سؤالات ساده و کوتاه، یک پاراگراف مستقیم کافی است.

## قوانین لحن
- لحن طبیعی، نه رسمی بیش از حد و نه خیلی غیررسمی.
- بدون جمله مقدمه اضافی. مستقیم پاسخ دهید.
- بدون سؤال فروشی در انتها. اگر اطلاعات تکمیلی مفید دارید، مستقیم بدهید.
- اگر کاربر "Momento" می‌گوید، منظورش Memento (این برنامه) است.`,
    },
    es: {
      contextWithNotes: `## Notas del usuario\n\n${contextNotes}\n\nCuando uses información de las notas anteriores, cita el título de la nota fuente entre paréntesis. No copies palabra por palabra — reformula. Si las notas no cubren el tema, dilo y complementa con tu conocimiento general.`,
      contextNoNotes: "No se encontraron notas relevantes para esta pregunta. Responde con tu conocimiento general.",
      system: `Eres el asistente de IA de Memento. El usuario te hace preguntas sobre sus proyectos, documentación técnica y notas. Debes responder de forma estructurada y útil.

## Reglas de formato
- Usa markdown libremente: títulos (##, ###), listas, bloques de código, negritas, tablas.
- Estructura tu respuesta con secciones para preguntas técnicas o temas complejos.
- Para preguntas simples y cortas, un párrafo directo es suficiente.

## Reglas de tono
- Tono natural, ni corporativo ni demasiado informal.
- Sin frases de introducción innecesarias. Responde directamente.
- Sin preguntas de venta al final. Si tienes información complementaria útil, dala directamente.`,
    },
    de: {
      contextWithNotes: `## Notizen des Benutzers\n\n${contextNotes}\n\nWenn du Infos aus den obigen Notizen verwendest, zitiere den Titel der Quellnotiz in Klammern. Nicht Wort für Wort kopieren — umformulieren. Wenn die Notizen das Thema nicht abdecken, sag es und ergänze mit deinem Allgemeinwissen.`,
      contextNoNotes: "Keine relevanten Notizen für diese Frage gefunden. Antworte mit deinem Allgemeinwissen.",
      system: `Du bist der KI-Assistent von Memento. Der Benutzer stellt dir Fragen zu seinen Projekten, technischen Dokumentationen und Notizen. Du musst strukturiert und hilfreich antworten.

## Formatregeln
- Verwende Markdown frei: Überschriften (##, ###), Listen, Code-Blöcke, Fettdruck, Tabellen.
- Strukturiere deine Antwort mit Abschnitten bei technischen Fragen oder komplexen Themen.
- Bei einfachen, kurzen Fragen reicht ein direkter Absatz.

## Tonregeln
- Natürlicher Ton, weder zu geschäftsmäßig noch zu umgangssprachlich.
- Keine unnötigen Einleitungssätze. Antworte direkt.
- Keine Upsell-Fragen am Ende. Gib nützliche Zusatzinfos einfach direkt.`,
    },
    it: {
      contextWithNotes: `## Note dell'utente\n\n${contextNotes}\n\nQuando usi informazioni dalle note sopra, cita il titolo della nota fonte tra parentesi. Non copiare parola per parola — riformula. Se le note non coprono l'argomento, dillo e integra con la tua conoscenza generale.`,
      contextNoNotes: "Nessuna nota rilevante trovata per questa domanda. Rispondi con la tua conoscenza generale.",
      system: `Sei l'assistente IA di Memento. L'utente ti fa domande sui suoi progetti, documentazione tecnica e note. Devi rispondere in modo strutturato e utile.

## Regole di formato
- Usa markdown liberamente: titoli (##, ###), elenchi, blocchi di codice, grassetto, tabelle.
- Struttura la risposta con sezioni per domande tecniche o argomenti complessi.
- Per domande semplici e brevi, un paragrafo diretto basta.

## Regole di tono
- Tono naturale, né aziendale né troppo informale.
- Nessuna frase introduttiva non necessaria. Rispondi direttamente.
- Nessuna domanda di upsell alla fine. Se hai informazioni aggiuntive utili, dalle direttamente.`,
    },
  }

  // Fallback to English if language not supported
  const prompts = promptLang[lang] || promptLang.en
  const contextBlock = contextNotes.length > 0
    ? prompts.contextWithNotes
    : prompts.contextNoNotes

  const systemPrompt = `${prompts.system}

${contextBlock}

${lang === 'en' ? 'Respond in the user\'s language.' : lang === 'fr' ? 'Réponds dans la langue de l\'utilisateur.' : 'Respond in the user\'s language.'}`

  // 6. Build message history from DB + current messages
  const dbHistory = conversation.messages.map((m: { role: string; content: string }) => ({
    role: m.role as 'user' | 'assistant' | 'system',
    content: m.content,
  }))

  // Only add the current user message if it's not already in DB history
  const lastIncoming = incomingMessages[incomingMessages.length - 1]
  const currentDbMessage = dbHistory[dbHistory.length - 1]
  const isNewMessage =
    lastIncoming &&
    (!currentDbMessage ||
      currentDbMessage.role !== 'user' ||
      currentDbMessage.content !== lastIncoming.content)

  const allMessages: Array<{ role: 'user' | 'assistant' | 'system'; content: string }> = isNewMessage
    ? [...dbHistory, { role: lastIncoming.role, content: lastIncoming.content }]
    : dbHistory

  // 7. Get chat provider model
  const config = await getSystemConfig()
  const provider = getChatProvider(config)
  const model = provider.getModel()

  // 8. Save user message to DB before streaming
  if (isNewMessage && lastIncoming) {
    await prisma.chatMessage.create({
      data: {
        conversationId: conversation.id,
        role: 'user',
        content: lastIncoming.content,
      },
    })
  }

  // 9. Stream response
  const result = streamText({
    model,
    system: systemPrompt,
    messages: allMessages,
    async onFinish({ text }) {
      // Save assistant message to DB after streaming completes
      await prisma.chatMessage.create({
        data: {
          conversationId: conversation.id,
          role: 'assistant',
          content: text,
        },
      })
    },
  })

  // 10. Return streaming response with conversation ID header
  return result.toUIMessageStreamResponse({
    headers: {
      'X-Conversation-Id': conversation.id,
    },
  })
}