Momento/memento-note/app/api/ai/suggest-charts/route.ts

import { generateText } from 'ai'
import { resolveAiRouteWithTiming } from '@/lib/ai/router'
import { runLaneWithBillingUser, willUseByokForLane } from '@/lib/ai/provider-for-user'
import { getSystemConfig } from '@/lib/config'
import { prisma } from '@/lib/prisma'
import { auth } from '@/auth'
import { toolRegistry } from '@/lib/ai/tools'
import { checkEntitlementOrThrow, QuotaExceededError } from '@/lib/entitlements'
import { trackFeatureUsage } from '@/lib/usage-tracker'

export const maxDuration = 30

interface SuggestChartsRequest {
  content: string
  selection?: string | null
  noteId?: string
}

interface ChartSuggestion {
  type: 'bar' | 'horizontal-bar' | 'line' | 'area' | 'pie' | 'radar' | 'funnel' | 'gauge'
  title: string
  data: { label: string; value: number }[]
  description: string
  rationale?: string
}

interface SuggestChartsResponse {
  suggestions: ChartSuggestion[]
  analyzedText: string
  detectedData: string
  hasData: boolean
}

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

  // 1.5 Quota check
  try {
    const sysConfigEarly = await getSystemConfig()
    const { usedByok: willUseByok } = await willUseByokForLane('suggest-charts', sysConfigEarly, userId)
    if (!willUseByok) {
      await checkEntitlementOrThrow(userId, 'ai')
    }
  } catch (err) {
    if (err instanceof QuotaExceededError) {
      return Response.json(err.toJSON(), { status: 402 })
    }
    console.error('[suggest-charts] Quota check error (fail-open):', err)
  }

  // 2. Parse request body
  const body = await req.json() as SuggestChartsRequest
  const { content, selection, noteId } = body

  if (!content || content.trim().length === 0) {
    return Response.json({
      error: 'Content is required',
      suggestions: [],
      analyzedText: '',
      detectedData: '',
      hasData: false,
    } satisfies SuggestChartsResponse, { status: 400 })
  }

  const textToAnalyze = selection && selection.trim() ? selection.trim() : content.trim()

  // 3. Build AI context
  const sysConfig = await getSystemConfig()
  const { provider, model, timingInfo } = await resolveAiRouteWithTiming('suggest-charts', sysConfig)

  // Build context for AI
  const toolContext = {
    userId,
    noteId,
    config: sysConfig,
  }

  // Use the suggest_charts tool
  const suggestChartsTool = toolRegistry.get('suggest_charts')
  const tools = suggestChartsTool ? { suggest_charts: suggestChartsTool.buildTool(toolContext) } : {}

  try {
    // 4. Call AI to analyze and suggest
    const { text } = await generateText({
      model: provider(model),
      system: `You are a data visualization assistant. Analyze the provided text and suggest appropriate chart types.

CRITICAL: Extract ONLY numerical data present in the text. Do NOT invent values.
- If fewer than 2 data points exist, return hasData=false with empty suggestions
- Each suggestion must use the SAME extracted data (only chart type differs)
- Return exactly 3 suggestions when data exists
- Provide clear rationale for each chart type choice

Response format (JSON):
{
  "suggestions": [
    {
      "type": "bar|horizontal-bar|line|area|pie|radar|funnel|gauge",
      "title": "Chart title",
      "data": [{"label": "Category", "value": 123}],
      "description": "What this chart shows",
      "rationale": "Why this chart type suits the data"
    }
  ],
  "analyzedText": "Text that was analyzed",
  "detectedData": "Description of what data was found",
  "hasData": true
}`,
      messages: [
        {
          role: 'user',
          content: `Analyze this text and suggest 3 appropriate chart types:\n\n${textToAnalyze}`,
        },
      ],
      tools,
      temperature: 0.3,
    })

    // 5. Parse AI response
    let parsed: SuggestChartsResponse
    try {
      // Try to extract JSON from the response
      const jsonMatch = text.match(/\{[\s\S]*\}/)
      if (!jsonMatch) {
        throw new Error('No JSON found in response')
      }
      parsed = JSON.parse(jsonMatch[0])
    } catch (e) {
      console.error('[suggest-charts] Failed to parse AI response:', text)
      // Return empty suggestions on parse error
      return Response.json({
        suggestions: [],
        analyzedText: textToAnalyze.substring(0, 100),
        detectedData: 'Failed to parse AI response',
        hasData: false,
      } satisfies SuggestChartsResponse)
    }

    // Validate response structure
    if (!parsed.suggestions || !Array.isArray(parsed.suggestions)) {
      parsed.suggestions = []
    }

    // Ensure exactly 3 suggestions when data exists
    if (parsed.hasData && parsed.suggestions.length !== 3) {
      parsed.suggestions = parsed.suggestions.slice(0, 3)
    }

    // Validate each suggestion has required fields
    parsed.suggestions = parsed.suggestions.filter(s =>
      s.type && s.title && s.data && Array.isArray(s.data) && s.data.length >= 2
    )

    // Track usage
    await trackFeatureUsage(userId, 'ai', 'suggest-charts', 1)

    return Response.json(parsed satisfies SuggestChartsResponse)

  } catch (error) {
    console.error('[suggest-charts] Error:', error)
    return Response.json({
      error: 'Failed to generate chart suggestions',
      suggestions: [],
      analyzedText: textToAnalyze.substring(0, 100),
      detectedData: 'Error occurred during analysis',
      hasData: false,
    } satisfies SuggestChartsResponse, { status: 500 })
  }
}