import { type ToolSet, type TypedToolCall, type TypedToolResult } from 'ai'

import { type AssistantEvalOutput } from './scorer'

type Step = {
  text: string
  toolCalls: TypedToolCall<ToolSet>[]
  toolResults: TypedToolResult<ToolSet>[]
}

type ParsedToolCall = {
  /** Query generated by `execute_sql` */
  sqlQuery?: string
  /** Docs text pulled in from `search_docs` */
  docs?: string[]
}

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) => {
          if (typeof text !== 'string') return false
          try {
            return !JSON.parse(text)?.error
          } catch {
            return true
          }
        })
      if (docs.length === 0) return {}
      return { docs }
    }
  }
  return {}
}

export function buildAssistantEvalOutput(
  finishReason: AssistantEvalOutput['finishReason'],
  steps: Step[]
): AssistantEvalOutput {
  const simplifiedSteps = steps.map((step) => ({
    text: step.text,
    toolCalls: step.toolCalls.map((call) => ({
      toolName: call.toolName,
      input: call.input,
    })),
  }))

  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 { finishReason, steps: simplifiedSteps, toolNames, sqlQueries, docs }
}