import OpenAI from 'openai'
import { createLogger } from '../../shared/logging/logger'
import { cleanManimCode } from '../../utils/manim-code-cleaner'
import { getClient } from './client'
import type { CodeRetryContext, RetryCheckpoint } from './types'
import { buildVisionUserMessage } from '../generation/multimodal-input'
import { buildRetryPrompt, getCodeRetrySystemPrompt } from './prompt-builder'
import { dedupeSharedBlocksInMessages } from '../prompt-dedup'
import { createChatCompletionText } from '../openai-stream'
import { buildCompletionParams, readMaxCompletionTokens } from '../../infrastructure/ai-provider/completion-policy'
import { applyPatchSetToCode, extractTargetLine } from './utils'
import { getApiRagService } from '../../knowledge/api-rag'
import { runRepairSession } from './repair-session'
import { inferAutomaticApiRequests } from './api-hints'

const logger = createLogger('CodeRetryCodeGen')

const AI_TEMPERATURE = parseFloat(process.env.AI_TEMPERATURE || '0.7')
const MAX_TOKENS = readMaxCompletionTokens()

function getModel(customApiConfig?: unknown): string {
  const model = (customApiConfig as { model?: string } | undefined)?.model
  const trimmed = model?.trim() || ''
  if (!trimmed) {
    throw new Error('No model available')
  }
  return trimmed
}

export async function retryCodeGeneration(
  context: CodeRetryContext,
  errorMessage: string,
  attempt: number,
  currentCode: string,
  codeSnippet: string | undefined,
  customApiConfig?: unknown,
  onCheckpoint?: RetryCheckpoint
): Promise<string> {
  const client = getClient(customApiConfig as any)
  if (!client) {
    throw new Error('No upstream AI is configured for this request')
  }

  const retryPrompt = buildRetryPrompt(context, errorMessage, attempt, currentCode, codeSnippet)
  const apiRagService = getApiRagService()

  try {
    const automaticRequests = inferAutomaticApiRequests(currentCode, errorMessage)
    const automaticResults = await Promise.all(
      automaticRequests.map((request) => apiRagService.search('manim', request))
    )
    const automaticApiContext = automaticResults.length > 0
      ? `\n\n[[AUTOMATIC_API_CONTEXT]]\n${automaticResults
          .map((result) => JSON.stringify(result, null, 2))
          .join('\n\n---\n\n')}\n[[END]]\nUse this verified API context. If it is insufficient, return an API_REQUEST.`
      : ''

    logger.info('Automatic retry API context resolved', {
      attempt,
      queries: automaticRequests.map((request) => request.query),
      candidates: automaticResults.map((result) => result.candidates.length)
    })

    const requestMessages: OpenAI.Chat.Completions.ChatCompletionMessageParam[] = dedupeSharedBlocksInMessages(
      [
        { role: 'system', content: getCodeRetrySystemPrompt(context.promptOverrides) },
        { role: 'user', content: `${retryPrompt}${automaticApiContext}` }
      ],
      context.promptOverrides
    ).map((message) => message.role === 'user'
      ? { role: 'user', content: buildVisionUserMessage(message.content, context.referenceImages) }
      : message)

    const session = await runRepairSession({
      messages: requestMessages,
      apiRagService,
      packageName: 'manim',
      concept: context.concept,
      attempt,
      onCheckpoint,
      requestTurn: (messages, turn) => createChatCompletionText(
        client,
        {
          model: getModel(customApiConfig),
          messages,
          temperature: AI_TEMPERATURE,
          ...buildCompletionParams(MAX_TOKENS, undefined, client)
        },
        { fallbackToNonStream: true, usageLabel: `retry-${attempt}-turn-${turn}` }
      )
    })

    const patchedCode = applyPatchSetToCode(
      currentCode,
      session.patchSet,
      extractTargetLine(errorMessage)
    )
    const cleaned = cleanManimCode(patchedCode)

    logger.info('Code retry patch applied', {
      concept: context.concept,
      attempt,
      turn: session.turn,
      mode: session.mode,
      patchCount: session.patchSet.patches.length,
      codeLength: cleaned.code.length,
      patchLengths: session.patchSet.patches.map((patch) => ({
        originalSnippetLength: patch.originalSnippet.length,
        replacementSnippetLength: patch.replacementSnippet.length
      })),
      codePreview: cleaned.code.slice(0, 500)
    })

    return cleaned.code
  } catch (error) {
    if (error instanceof OpenAI.APIError) {
      logger.error('OpenAI API error during code retry', {
        attempt,
        status: error.status,
        message: error.message
      })
    }
    throw error
  }
}