import { openai } from '@ai-sdk/openai'
import { Eval } from 'braintrust'
import { generateAssistantResponse } from 'lib/ai/generate-assistant-response'
import { getMockTools } from 'lib/ai/tools/mock-tools'
import assert from 'node:assert'
import { dataset } from './dataset'
import {
  completenessScorer,
  concisenessScorer,
  goalCompletionScorer,
  sqlSyntaxScorer,
  toolUsageScorer,
} from './scorer'
import { ToolSet, TypedToolCall, TypedToolResult } from 'ai'

assert(process.env.BRAINTRUST_PROJECT_ID, 'BRAINTRUST_PROJECT_ID is not set')
assert(process.env.OPENAI_API_KEY, 'OPENAI_API_KEY is not set')

Eval('Assistant', {
  projectId: process.env.BRAINTRUST_PROJECT_ID,
  data: () => dataset,
  task: async (input) => {
    const result = await generateAssistantResponse({
      model: openai('gpt-5-mini'),
      messages: [{ id: '1', role: 'user', parts: [{ type: 'text', text: input }] }],
      tools: await getMockTools(),
    })

    // `result.toolCalls` only shows the last step, instead aggregate tools across all steps
    const steps = await result.steps

    const stepsSerialized = steps
      .map((step) => {
        const toolCalls = step.toolCalls
          ?.map((call) => JSON.stringify({ tool: call.toolName, input: call.input }))
          .join('\n')

        const text = step.text
        return toolCalls ? `${text}\n${toolCalls}` : text
      })
      .join('\n')

    const toolNames: string[] = []
    const sqlQueries: string[] = []
    const docs: string[] = []

    for (const step of steps) {
      for (const [i, toolCall] of step.toolCalls.entries()) {
        toolNames.push(toolCall.toolName)

        const toolResult = step.toolResults.at(i)
        if (!toolResult) {
          continue
        }

        const parsed = parseToolCall(toolCall, toolResult)

        if (parsed.sqlQuery) {
          sqlQueries.push(parsed.sqlQuery)
        }
        if (parsed.docs) {
          docs.push(...parsed.docs)
        }
      }
    }

    return {
      stepsSerialized,
      toolNames,
      sqlQueries,
      docs,
    }
  },
  scores: [
    toolUsageScorer,
    sqlSyntaxScorer,
    goalCompletionScorer,
    concisenessScorer,
    completenessScorer,
  ],
})

type ParsedToolCall = {
  /** Query generated by `execute_sql` */
  sqlQuery?: string

  /** Docs text pulled in from `search_docs` */
  docs?: string[]
}

/**
 * Validate and extract relevant info from a tool call/result
 */
function parseToolCall(
  toolCall: TypedToolCall<ToolSet>,
  toolResult: TypedToolResult<ToolSet>
): ParsedToolCall {
  switch (toolCall.toolName) {
    case 'execute_sql': {
      const sqlQuery = toolCall.input.sql
      if (typeof sqlQuery !== 'string') {
        return {}
      }

      return { sqlQuery }
    }
    case 'search_docs': {
      const content = toolResult.output.content
      if (!content || !Array.isArray(content)) {
        return {}
      }

      const docs = content.map((item) => item?.text).filter((text) => typeof text === 'string')
      if (docs.length === 0) {
        return {}
      }

      return { docs }
    }
  }

  return {}
}