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, 'suggest_charts')
    }
  } 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)

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

IMPORTANT: You MUST respond with valid JSON only. No markdown, no code blocks, no explanations.

DATA EXTRACTION RULES:
- Find ANY numbers in the text - even simple lists count
- Look for: "X: 123", "X = 123", "123", "123%", "$123", etc.
- If you find ANY 2+ numbers, create a chart
- Be creative with labels if none exist (use "Item 1", "Item 2", etc.)
- NEVER return hasData=false unless text is completely empty or has no numbers at all

CHART TYPES TO SUGGEST (always 3 different types):
1. bar - for comparing categories
2. line - for trends/sequences
3. pie - for parts of whole

Response format (COPY this structure):
{"suggestions":[
  {"type":"bar","title":"Chart Title","data":[{"label":"A","value":100},{"label":"B","value":200}],"description":"...","rationale":"..."},
  {"type":"line","title":"...","data":[...],"description":"...","rationale":"..."},
  {"type":"pie","title":"...","data":[...],"description":"...","rationale":"..."}
],"analyzedText":"...","detectedData":"...","hasData":true}`,
      messages: [
        {
          role: 'user',
          content: `Extract numbers from this text and suggest 3 charts:\n\n${textToAnalyze}`,
        },
      ],
      temperature: 0.3,
    })

    // 5. Parse AI response - be very lenient
    console.log('[suggest-charts] AI response:', text.substring(0, 500))

    let parsed: SuggestChartsResponse
    try {
      // Clean the response - remove markdown code blocks
      let cleanText = text
        .replace(/```json\n?/gi, '')
        .replace(/```\n?/gi, '')
        .trim()

      // Find JSON object
      const jsonMatch = cleanText.match(/\{[\s\S]*\}/)
      if (!jsonMatch) {
        throw new Error('No JSON found')
      }
      parsed = JSON.parse(jsonMatch[0])
    } catch (e) {
      console.error('[suggest-charts] Parse error:', e, 'Raw text:', text)

      // Check if text has ANY numbers - if so, create a simple chart
      const numbers = textToAnalyze.match(/\d+/g)
      if (numbers && numbers.length >= 2) {
        const values = numbers.slice(0, 6).map(n => parseInt(n) || 0)
        parsed = {
          suggestions: [
            {
              type: 'bar',
              title: 'Data from Note',
              data: values.map((v, i) => ({ label: `Item ${i + 1}`, value: v })),
              description: 'Bar chart of extracted values',
              rationale: 'Simple comparison of numerical values found'
            },
            {
              type: 'line',
              title: 'Data Trend',
              data: values.map((v, i) => ({ label: `Item ${i + 1}`, value: v })),
              description: 'Line chart showing progression',
              rationale: 'Visualizes the sequence of values'
            },
            {
              type: 'pie',
              title: 'Data Distribution',
              data: values.map((v, i) => ({ label: `Item ${i + 1}`, value: v })),
              description: 'Pie chart of value proportions',
              rationale: 'Shows relative sizes of each value'
            }
          ],
          analyzedText: textToAnalyze.substring(0, 100),
          detectedData: `Found ${numbers.length} numerical values`,
          hasData: true
        }
      } else {
        return Response.json({
          suggestions: [],
          analyzedText: textToAnalyze.substring(0, 100),
          detectedData: 'No numerical data found',
          hasData: false,
          error: 'Add numbers to your note (e.g., "Jan: 100, Feb: 200")',
        } satisfies SuggestChartsResponse, { status: 200 })
      }
    }

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

    // Ensure we have 3 suggestions when hasData=true
    if (parsed.hasData && parsed.suggestions.length < 3) {
      const baseSuggestion = parsed.suggestions[0]
      if (baseSuggestion) {
        const types = ['bar', 'line', 'pie'].filter(t => t !== baseSuggestion.type)
        while (parsed.suggestions.length < 3 && types.length > 0) {
          parsed.suggestions.push({ ...baseSuggestion, type: types.shift()! })
        }
      }
    }

    // Ensure each suggestion has valid data
    parsed.suggestions = parsed.suggestions.filter(s => {
      if (!s.type || !s.data || !Array.isArray(s.data) || s.data.length < 2) return false
      // Ensure all data points have label and value
      s.data = s.data.filter(d => d && typeof d.value === 'number')
      return s.data.length >= 2
    })

    // If after filtering we have no valid suggestions, set hasData=false
    if (parsed.suggestions.length === 0) {
      parsed.hasData = false
    }

    // Track usage
    await trackFeatureUsage(userId, 'suggest_charts', '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 })
  }
}