src/services/api/nvidiaClient.ts

import { getNvidiaApiKey } from '../../utils/auth.js'
import { getNvidiaBaseUrl } from '../../utils/model/providers.js'
import {
  convertAnthropicMessagesToOpenAI,
  convertAnthropicToolsToOpenAI,
  convertOpenAIStreamToAnthropic,
  type AnthropicMessage,
} from './copilotClient.js'

/**
 * NVIDIA NIM API uses the OpenAI-compatible `/v1/chat/completions` protocol.
 * This fetch override intercepts Anthropic Messages API calls and translates
 * them to the OpenAI format that NVIDIA expects.
 */

const NVIDIA_MODEL = process.env.NVIDIA_MODEL || 'nvidia/llama-3.1-nemotron-70b-instruct'

function normalizeBaseUrl(url: string): string {
  return url.replace(/\/$/, '')
}

function chatCompletionsUrl(base: string): string {
  const b = normalizeBaseUrl(base)
  if (b.endsWith('/v1')) {
    return `${b}/chat/completions`
  }
  return `${b}/v1/chat/completions`
}

export function createNvidiaFetchOverride(): (input: RequestInfo | URL, init?: RequestInit) => Promise<Response> {
  const baseUrl = getNvidiaBaseUrl()
  const apiKey = getNvidiaApiKey()
  const endpoint = chatCompletionsUrl(baseUrl)

  return async (input: RequestInfo | URL, init?: RequestInit): Promise<Response> => {
    const url = input instanceof URL ? input.href : typeof input === 'string' ? input : input.url

    // Only intercept Messages API calls
    if (!url.includes('/messages') && !url.includes('/v1/')) {
      return fetch(input, init)
    }

    // Stub out token counting and model listing
    if (url.includes('/count_tokens') || url.includes('/models')) {
      return new Response(JSON.stringify({ input_tokens: 0 }), {
        status: 200,
        headers: { 'Content-Type': 'application/json' },
      })
    }

    let anthropicBody: Record<string, unknown> = {}
    if (init?.body) {
      try {
        anthropicBody = JSON.parse(
          typeof init.body === 'string' ? init.body : new TextDecoder().decode(init.body as ArrayBuffer),
        )
      } catch {
        return fetch(input, init)
      }
    }

    const systemBlocks = anthropicBody.system as
      | Array<{ type: string; text: string }>
      | string
      | undefined
    let systemPrompt = ''
    if (typeof systemBlocks === 'string') {
      systemPrompt = systemBlocks
    } else if (Array.isArray(systemBlocks)) {
      systemPrompt = systemBlocks
        .filter(b => b.type === 'text')
        .map(b => b.text)
        .join('\n\n')
    }

    const anthropicMessages = (anthropicBody.messages || []) as AnthropicMessage[]
    const openaiMessages = convertAnthropicMessagesToOpenAI(anthropicMessages, systemPrompt)

    const anthropicTools = (anthropicBody.tools || []) as Array<{
      name: string
      description?: string
      input_schema?: Record<string, unknown>
    }>
    const openaiTools = anthropicTools.length > 0 ? convertAnthropicToolsToOpenAI(anthropicTools) : undefined

    const isStreaming = anthropicBody.stream === true

    // Use the model from the user's selection, falling back to the env var / default
    const selectedModel = (anthropicBody.model as string) || NVIDIA_MODEL

    const requestBody: Record<string, unknown> = {
      model: selectedModel,
      messages: openaiMessages,
      stream: isStreaming,
    }

    if (anthropicBody.max_tokens) {
      requestBody.max_tokens = anthropicBody.max_tokens
    }

    if (openaiTools && openaiTools.length > 0) {
      requestBody.tools = openaiTools
      requestBody.tool_choice = 'auto'
    }

    const headers: Record<string, string> = {
      'Content-Type': 'application/json',
      'User-Agent': 'claude-code/2.1.88',
      'HTTP-Referer': 'https://claude.ai/',
      'X-Title': 'Better-Clawd',
      'X-BILLING-INVOKE-ORIGIN': 'Better-Clawd',
    }
    if (apiKey) {
      headers.Authorization = `Bearer ${apiKey}`
    }

    const nvidiaResponse = await fetch(endpoint, {
      method: 'POST',
      headers,
      body: JSON.stringify(requestBody),
      signal: init?.signal,
    })

    if (!nvidiaResponse.ok) {
      return nvidiaResponse
    }

    if (!isStreaming) {
      const data = (await nvidiaResponse.json()) as {
        id: string
        choices: Array<{
          message: {
            role: string
            content: string | null
            tool_calls?: Array<{
              id: string
              function: { name: string; arguments: string }
            }>
          }
          finish_reason: string
        }>
        usage?: { prompt_tokens: number; completion_tokens: number }
      }

      const choice = data.choices[0]
      const anthropicContent: Array<{
        type: string
        text?: string
        id?: string
        name?: string
        input?: unknown
      }> = []

      if (choice?.message?.content) {
        anthropicContent.push({ type: 'text', text: choice.message.content })
      }

      if (choice?.message?.tool_calls) {
        for (const tc of choice.message.tool_calls) {
          anthropicContent.push({
            type: 'tool_use',
            id: tc.id,
            name: tc.function.name,
            input: JSON.parse(tc.function.arguments || '{}'),
          })
        }
      }

      const anthropicResponse = {
        id: data.id || `msg_nvidia_${Date.now()}`,
        type: 'message',
        role: 'assistant',
        content: anthropicContent,
        model: selectedModel,
        stop_reason: choice?.finish_reason === 'tool_calls' ? 'tool_use' : 'end_turn',
        usage: {
          input_tokens: data.usage?.prompt_tokens || 0,
          output_tokens: data.usage?.completion_tokens || 0,
        },
      }

      return new Response(JSON.stringify(anthropicResponse), {
        status: 200,
        headers: { 'Content-Type': 'application/json' },
      })
    }

    // Streaming response
    if (!nvidiaResponse.body) {
      return nvidiaResponse
    }

    const transformStream = convertOpenAIStreamToAnthropic(nvidiaResponse.body, selectedModel)

    return new Response(transformStream, {
      status: 200,
      headers: {
        'Content-Type': 'text/event-stream',
        'Cache-Control': 'no-cache',
        Connection: 'keep-alive',
      },
    })
  }
}

let cachedNvidiaModels: string[] | null = null

export function getCachedNvidiaModels(): string[] {
  return cachedNvidiaModels || []
}

/**
 * Fetch available models from the NVIDIA API catalog.
 */
export async function fetchNvidiaModels(apiKey?: string): Promise<string[]> {
  const baseUrl = getNvidiaBaseUrl()
  const key = apiKey || getNvidiaApiKey()
  const normalizedBase = normalizeBaseUrl(baseUrl)
  const modelsUrl = normalizedBase.endsWith('/v1')
    ? `${normalizedBase}/models`
    : `${normalizedBase}/v1/models`

  const headers: Record<string, string> = {}
  if (key) {
    headers.Authorization = `Bearer ${key}`
  }

  try {
    const res = await fetch(modelsUrl, { headers, signal: AbortSignal.timeout(20_000) })
    if (!res.ok) {
      cachedNvidiaModels = []
      return []
    }
    const json = (await res.json()) as { data?: Array<{ id: string }> }
    if (json.data && Array.isArray(json.data)) {
      const modelIds = json.data.map((m: { id: string }) => m.id)
      cachedNvidiaModels = modelIds
      return modelIds
    }
    cachedNvidiaModels = []
    return []
  } catch {
    cachedNvidiaModels = []
    return []
  }
}