Momento/memento-note/app/api/brainstorm/route.ts

import { NextRequest, NextResponse } from 'next/server'
import prisma from '@/lib/prisma'
import { auth } from '@/auth'
import { z } from 'zod'
import { runLaneWithBillingUser, willUseByokForLane } from '@/lib/ai/provider-for-user'
import { getSystemConfig } from '@/lib/config'
import { embeddingService } from '@/lib/ai/services/embedding.service'
import {
  reserveUsageOrThrow,
  QuotaExceededError,
} from '@/lib/entitlements'
import { logActivity, captureSnapshot } from '@/lib/brainstorm-collab'

const waveSchema = z.object({
  seedIdea: z.string().min(1, 'Seed idea is required'),
  sourceNoteId: z.string().optional(),
  contextNoteIds: z.array(z.string()).optional(),
  locale: z.string().optional(),
})

interface ClassifiedNote {
  id: string
  title: string
  summary: string
  category: 'SUPPORT' | 'TENSION' | 'EXTENSION'
}

async function autoContextSearch(
  userId: string,
  seedIdea: string,
  userNoteIds?: string[]
): Promise<ClassifiedNote[]> {
  let candidateIds: string[] = []

  if (userNoteIds && userNoteIds.length > 0) {
    candidateIds = userNoteIds
  } else {
    try {
      const embedding = await embeddingService.generateEmbedding(seedIdea)
      const vectorStr = embeddingService.toVectorString(embedding.embedding)
      const results = await prisma.$queryRawUnsafe(
        `SELECT n.id
         FROM "NoteEmbedding" e
         JOIN "Note" n ON n.id = e."noteId"
         WHERE n."userId" = $1 AND n."trashedAt" IS NULL
          ORDER BY e.embedding::vector <=> $2::vector
         LIMIT 8`,
        userId, vectorStr
      ) as any[]
      candidateIds = results.map((r: any) => r.id)
    } catch {
      return []
    }
  }

  if (candidateIds.length === 0) return []

  const notes = await prisma.note.findMany({
    where: { id: { in: candidateIds }, userId, trashedAt: null },
    select: { id: true, title: true, content: true },
  })

  if (notes.length === 0) return []

  const notesForLLM = notes.map(n => ({
    id: n.id,
    title: n.title || 'Untitled',
    snippet: (n.content || '').slice(0, 300),
  }))

  const classifyPrompt = `Given the seed idea: "${seedIdea}"

Classify each note as SUPPORT (confirms/reinforces the seed), TENSION (contradicts/questions the seed), or EXTENSION (extends the seed into an adjacent domain).

Notes:
${notesForLLM.map(n => `[${n.id}] "${n.title}": ${n.snippet}`).join('\n')}

Respond ONLY with a valid JSON array of objects:
{ "noteId": string, "category": "SUPPORT" | "TENSION" | "EXTENSION" }`

  try {
    const config = await getSystemConfig()
    const { result: raw } = await runLaneWithBillingUser(
      'tags',
      config,
      userId,
      (provider) => provider.generateText(classifyPrompt),
    )
    const cleaned = raw.replace(/```json\n?/g, '').replace(/```\n?/g, '').trim()
    const classifications: { noteId: string; category: 'SUPPORT' | 'TENSION' | 'EXTENSION' }[] = JSON.parse(cleaned)

    return notes.map(n => {
      const cls = classifications.find(c => c.noteId === n.id)
      return {
        id: n.id,
        title: n.title || 'Untitled',
        summary: (n.content || '').slice(0, 200),
        category: cls?.category || 'EXTENSION',
      }
    })
  } catch {
    return notes.map(n => ({
      id: n.id,
      title: n.title || 'Untitled',
      summary: (n.content || '').slice(0, 200),
      category: 'EXTENSION' as const,
    }))
  }
}

function localeToLanguageName(locale?: string): string {
  const map: Record<string, string> = {
    en: 'English', fr: 'French', es: 'Spanish', de: 'German',
    it: 'Italian', pt: 'Portuguese', nl: 'Dutch', ru: 'Russian',
    zh: 'Chinese', ja: 'Japanese', ko: 'Korean', ar: 'Arabic',
    fa: 'Farsi', hi: 'Hindi', pl: 'Polish',
  }
  return map[locale ?? ''] ?? 'English'
}

function buildPromptV2(seedIdea: string, classifiedNotes: ClassifiedNote[], locale?: string): string {
  const supportNotes = classifiedNotes.filter(n => n.category === 'SUPPORT')
  const tensionNotes = classifiedNotes.filter(n => n.category === 'TENSION')
  const extensionNotes = classifiedNotes.filter(n => n.category === 'EXTENSION')

  let notesSection = ''
  if (classifiedNotes.length > 0) {
    notesSection = `\nUSER'S EXISTING NOTES (classified by relationship to the seed):\n`
    if (supportNotes.length > 0) {
      notesSection += `\nSUPPORTING NOTES (confirm/reinforce the seed):\n`
      notesSection += supportNotes.map(n => `- [ID: ${n.id}] "${n.title}": ${n.summary}`).join('\n')
    }
    if (tensionNotes.length > 0) {
      notesSection += `\nTENSION NOTES (contradict or question the seed):\n`
      notesSection += tensionNotes.map(n => `- [ID: ${n.id}] "${n.title}": ${n.summary}`).join('\n')
    }
    if (extensionNotes.length > 0) {
      notesSection += `\nEXTENSION NOTES (extend the seed into adjacent domains):\n`
      notesSection += extensionNotes.map(n => `- [ID: ${n.id}] "${n.title}": ${n.summary}`).join('\n')
    }
  }

  return `You are a creative brainstorming assistant with access to the user's personal knowledge base. Your job is to generate ideas that DELIBERATELY CROSS the seed concept with existing notes — creating productive tension, not just variations.

USER'S SEED IDEA: ${seedIdea}
${notesSection}

GENERATION RULES:

WAVE 1 — VARIATIONS (3 ideas):
- At least 1 idea must BUILD ON a SUPPORTING note
- At least 1 idea must RESPOND TO a TENSION note (resolve or embrace the contradiction)

WAVE 2 — ANALOGIES (3 ideas):
- At least 1 idea must FUSE a concept from an EXTENSION note with the seed
- At least 1 idea must be a PATTERN TRANSPOSITION from any note

WAVE 3 — DISRUPTIONS (3 ideas):
- At least 1 idea must INVERT an assumption found in a SUPPORTING note
- At least 1 idea must SYNTHESIZE two notes that appear contradictory

RESPOND ONLY with a valid JSON array of 9 objects:
{
  "wave": number (1, 2, or 3),
  "title": string (short, 2-6 words),
  "description": string (1-2 sentences, specific and actionable),
  "connectionToSeed": string (how it relates to the seed),
  "noveltyScore": number (1-10),
  "noteRefs": [
    {
      "noteId": string (must match an ID provided above, or null if genuinely no note connects),
      "relation": "derived_from" | "opposes" | "extends" | "synthesizes" | "transposes",
      "explanation": string (natural language, e.g. "Built on your exploration of occupancy-based scheduling")
    }
  ]
}

  CRITICAL: Each idea MUST have at least 1 noteRef. Only use null noteId if genuinely no note connects to that idea.

LANGUAGE: You MUST write ALL titles, descriptions, connectionToSeed, and explanation fields in ${localeToLanguageName(locale)}.`
}

function buildFallbackIdeas(classifiedNotes: ClassifiedNote[]): any[] {
  const notesById = Object.fromEntries(classifiedNotes.map(n => [n.id, n]))
  const support = classifiedNotes.filter(n => n.category === 'SUPPORT')
  const tension = classifiedNotes.filter(n => n.category === 'TENSION')
  const extension = classifiedNotes.filter(n => n.category === 'EXTENSION')

  const pickNote = (list: ClassifiedNote[], idx: number) => {
    if (list.length === 0) return { noteId: null, relation: 'extends' as const, explanation: 'Purely generative idea, no direct note link' }
    const n = list[idx % list.length]
    return { noteId: n.id, relation: 'extends' as const, explanation: `Inspired by your note "${n.title}"` }
  }

  return [
    { wave: 1, title: 'Variation A', description: 'A direct variation of the seed idea.', connectionToSeed: 'Direct extension', noveltyScore: 3, noteRefs: [pickNote(support, 0)] },
    { wave: 1, title: 'Variation B', description: 'Another angle on the seed idea.', connectionToSeed: 'Reformulation', noveltyScore: 4, noteRefs: [pickNote(tension, 0)] },
    { wave: 1, title: 'Variation C', description: 'A sub-aspect of the seed.', connectionToSeed: 'Sub-component', noveltyScore: 5, noteRefs: [pickNote(support, 1)] },
    { wave: 2, title: 'Analogy A', description: 'Inspired by biological systems.', connectionToSeed: 'Cross-domain analogy', noveltyScore: 6, noteRefs: [pickNote(extension, 0)] },
    { wave: 2, title: 'Analogy B', description: 'Drawn from technology patterns.', connectionToSeed: 'Tech parallel', noveltyScore: 7, noteRefs: [pickNote(extension, 1)] },
    { wave: 2, title: 'Analogy C', description: 'Based on natural phenomena.', connectionToSeed: 'Nature metaphor', noveltyScore: 6, noteRefs: [pickNote(classifiedNotes, 0)] },
    { wave: 3, title: 'Disruption A', description: 'What if we inverted the core assumption?', connectionToSeed: 'Inversion', noveltyScore: 9, noteRefs: [{ ...(pickNote(support, 0)), relation: 'opposes' as const }] },
    { wave: 3, title: 'Disruption B', description: 'A provocative reframing.', connectionToSeed: 'Challenge premise', noveltyScore: 8, noteRefs: [{ ...(pickNote(tension, 0)), relation: 'synthesizes' as const }] },
    { wave: 3, title: 'Disruption C', description: 'Removing a key constraint entirely.', connectionToSeed: 'Constraint removal', noveltyScore: 10, noteRefs: [pickNote(extension, 2)] },
  ]
}

export async function POST(request: NextRequest) {
  const session = await auth()
  if (!session?.user?.id) {
    return NextResponse.json({ error: 'Unauthorized' }, { status: 401 })
  }
  const userId = session.user.id

  try {
    const body = await request.json()
    const { seedIdea, sourceNoteId, contextNoteIds, locale } = waveSchema.parse(body)

    // Story 3.5: per-provider BYOK bypass
    const config = await getSystemConfig()
    const { usedByok: willUseByok } = await willUseByokForLane('tags', config, userId)
    if (!willUseByok) {
      await reserveUsageOrThrow(userId, 'brainstorm_create')
    }

    const classifiedNotes = await autoContextSearch(userId, seedIdea, contextNoteIds)

    const prompt = buildPromptV2(seedIdea, classifiedNotes, locale)
    const { result: llmResponse, usedByok } = await runLaneWithBillingUser(
      'tags',
      config,
      userId,
      (provider) => provider.generateText(prompt),
    )

    let ideas: any[]
    try {
      const cleaned = llmResponse.replace(/```json\n?/g, '').replace(/```\n?/g, '').trim()
      ideas = JSON.parse(cleaned)
      if (!Array.isArray(ideas)) throw new Error('Not an array')
    } catch {
      ideas = buildFallbackIdeas(classifiedNotes)
    }

    const validNoteIds = new Set(classifiedNotes.map(n => n.id))

    const brainstormSession = await prisma.brainstormSession.create({
      data: {
        seedIdea,
        sourceNoteId: sourceNoteId || null,
        contextNoteIds: contextNoteIds ? JSON.stringify(contextNoteIds) : null,
        liveblocksRoomId: `brainstorm-session`,
        userId,
      },
    })

    await prisma.brainstormParticipant.create({
      data: {
        sessionId: brainstormSession.id,
        userId,
        role: 'host',
      },
    })

    await logActivity(brainstormSession.id, 'wave_generated', userId, { count: ideas.length })

    const createdIdeas = []
    for (let idx = 0; idx < ideas.length; idx++) {
      const idea = ideas[idx]
      const angle = (idx % 3) * (2 * Math.PI / 3) + (idea.wave - 1) * 0.5
      const radius = idea.wave * 150

      const created = await prisma.brainstormIdea.create({
        data: {
          sessionId: brainstormSession.id,
          waveNumber: idea.wave || Math.floor(idx / 3) + 1,
          title: idea.title || `Idea ${idx + 1}`,
          description: idea.description || '',
          connectionToSeed: idea.connectionToSeed || null,
          noveltyScore: idea.noveltyScore || null,
          relatedNoteIds: JSON.stringify(
            (idea.noteRefs || []).map((r: any) => r.noteId).filter(Boolean)
          ),
          positionX: Math.cos(angle) * radius,
          positionY: Math.sin(angle) * radius,
        },
      })

      if (idea.noteRefs && Array.isArray(idea.noteRefs)) {
        for (const ref of idea.noteRefs) {
          const noteId = ref.noteId && validNoteIds.has(ref.noteId) ? ref.noteId : null
          await prisma.brainstormNoteRef.create({
            data: {
              ideaId: created.id,
              noteId,
              relation: ref.relation || 'extends',
              explanation: ref.explanation || '',
            },
          })
        }
      }

      createdIdeas.push(created)
    }

    const fullSession = await prisma.brainstormSession.findUnique({
      where: { id: brainstormSession.id },
      include: {
        ideas: {
          orderBy: [{ waveNumber: 'asc' }, { createdAt: 'asc' }],
          include: {
            noteRefs: {
              include: {
                note: { select: { id: true, title: true } },
              },
            },
          },
        },
      },
    })

    const cIds = [...new Set((fullSession?.ideas || []).map((i: any) => i.createdBy).filter(Boolean))]
    if (cIds.length > 0) {
      const crs = await prisma.user.findMany({ where: { id: { in: cIds } }, select: { id: true, name: true, image: true } })
      const cm = new Map(crs.map((c: any) => [c.id, c]))
      for (const idea of fullSession?.ideas || []) { (idea as any).creator = (idea as any).createdBy ? cm.get((idea as any).createdBy) || null : null }
    }

    await captureSnapshot(brainstormSession.id, `Initial wave: ${brainstormSession.seedIdea.substring(0, 30)}`).catch(() => {})

    return NextResponse.json({
      success: true,
      data: fullSession,
      contextSummary: {
        support: classifiedNotes.filter(n => n.category === 'SUPPORT').length,
        tension: classifiedNotes.filter(n => n.category === 'TENSION').length,
        extension: classifiedNotes.filter(n => n.category === 'EXTENSION').length,
      },
    }, { status: 201 })
  } catch (error: any) {
    if (error instanceof QuotaExceededError) {
      return NextResponse.json(error.toJSON(), { status: 402 })
    }
    if (error instanceof z.ZodError) {
      return NextResponse.json({ error: error.issues }, { status: 400 })
    }
    console.error('Error creating brainstorm:', error)
    return NextResponse.json(
      { error: error.message || 'Failed to create brainstorm' },
      { status: 500 }
    )
  }
}

export async function GET(request: NextRequest) {
  const session = await auth()
  if (!session?.user?.id) {
    return NextResponse.json({ error: 'Unauthorized' }, { status: 401 })
  }
  const userId = session.user.id

  try {
    const ownedSessions = await prisma.brainstormSession.findMany({
      where: { userId },
      orderBy: { createdAt: 'desc' },
      include: {
        _count: { select: { ideas: true } },
        ideas: {
          where: { status: 'active' },
          select: { id: true },
        },
      },
    })

    const owned = ownedSessions.map(s => ({
      id: s.id,
      seedIdea: s.seedIdea,
      sourceNoteId: s.sourceNoteId,
      exportedNoteId: s.exportedNoteId,
      createdAt: s.createdAt,
      updatedAt: s.updatedAt,
      totalIdeas: s._count.ideas,
      activeIdeas: s.ideas.length,
      _owned: true,
    }))

    const ownedIds = new Set(owned.map(s => s.id))

    let shared: any[] = []
    try {
      const acceptedShareRows = await prisma.brainstormShare.findMany({
        where: { userId, status: 'accepted' },
        include: {
          session: {
            select: {
              id: true,
              seedIdea: true,
              sourceNoteId: true,
              exportedNoteId: true,
              createdAt: true,
              updatedAt: true,
              _count: { select: { ideas: true } },
              ideas: { where: { status: 'active' }, select: { id: true } },
            },
          },
        },
      })
      shared = acceptedShareRows
        .filter(s => s.session != null && !ownedIds.has(s.session.id))
        .map(s => ({
          id: s.session.id,
          seedIdea: s.session.seedIdea,
          sourceNoteId: s.session.sourceNoteId,
          exportedNoteId: s.session.exportedNoteId,
          createdAt: s.session.createdAt,
          updatedAt: s.session.updatedAt,
          totalIdeas: s.session._count.ideas,
          activeIdeas: s.session.ideas.length,
          _owned: false,
        }))
    } catch (shareError) {
      console.error('Error fetching shared brainstorms (non-fatal):', shareError)
    }

    return NextResponse.json({
      success: true,
      data: [...owned, ...shared],
    })
  } catch (error) {
    console.error('Error fetching brainstorms:', error)
    return NextResponse.json(
      { error: 'Failed to fetch brainstorms' },
      { status: 500 }
    )
  }
}