import { createRetryContext, executeCodeRetry } from '../../../../workflow/code-retry/manager'
import type { ConceptInput, OutputMode, PromptOverrides, ProblemFramingPlan } from '../../../../types'

export interface RenderAttemptResult {
  success: boolean
  stderr: string
  stdout: string
  peakMemoryMB: number
  exitCode?: number
  codeSnippet?: string
}

interface RenderFailureLogArgs {
  attempt: number
  code: string
  codeSnippet?: string
  stderr: string
  stdout: string
  peakMemoryMB: number
  exitCode?: number
  promptRole: 'codeRetry' | 'single-render'
}

interface ExecuteRenderWithRetryArgs extends ConceptInput {
  problemPlan?: ProblemFramingPlan
  outputMode: OutputMode
  sceneDesign?: string
  promptOverrides?: PromptOverrides
  customApiConfig?: unknown
  initialCode: string
  usedAI: boolean
  timings?: Record<string, number>
  renderCode: (code: string) => Promise<RenderAttemptResult>
  logRenderFailure: (args: RenderFailureLogArgs) => Promise<void>
  ensureJobNotCancelled?: () => Promise<void>
  onRetryStart?: () => Promise<void> | void
}

export async function executeRenderWithRetry(args: ExecuteRenderWithRetryArgs): Promise<{
  finalCode: string
  renderPeakMemoryMB: number
}> {
  const {
    concept,
    outputMode,
    sceneDesign,
    promptOverrides,
    customApiConfig,
    initialCode,
    usedAI,
    timings,
    renderCode,
    logRenderFailure,
    ensureJobNotCancelled,
    onRetryStart
  } = args

  if (usedAI) {
    const retryContext = createRetryContext(
      concept,
      sceneDesign?.trim() || `概念: ${concept}`,
      promptOverrides,
      outputMode,
      args.problemPlan,
      args.referenceImages
    )

    const retryManagerResult = await executeCodeRetry(
      retryContext,
      renderCode,
      customApiConfig,
      initialCode,
      async (event) => {
        await logRenderFailure({ ...event, promptRole: 'codeRetry' })
      },
      ensureJobNotCancelled,
      timings,
      onRetryStart
    )

    if (typeof retryManagerResult.generationTimeMs === 'number' && timings) {
      timings.retry = retryManagerResult.generationTimeMs
    }

    if (!retryManagerResult.success) {
      // A cancelled render is not a code failure. Re-raise the job's cancellation with its
      // original JobCancelledError instead of the retry manager's wrapped diagnostic.
      if (retryManagerResult.failureCategory === 'cancelled') {
        await ensureJobNotCancelled?.()
      }
      const failurePrefix = retryManagerResult.retrySkipped
        ? `AI retry skipped for ${retryManagerResult.failureCategory || 'operational'} failure`
        : `Code retry failed after ${retryManagerResult.attempts} attempts`
      throw new Error(
        `${failurePrefix}: ${retryManagerResult.lastError}`
      )
    }

    return {
      finalCode: retryManagerResult.code,
      renderPeakMemoryMB: 0
    }
  }

  const singleAttempt = await renderCode(initialCode)
  if (!singleAttempt.success) {
    await logRenderFailure({
      attempt: 1,
      code: initialCode,
      codeSnippet: singleAttempt.codeSnippet,
      stderr: singleAttempt.stderr,
      stdout: singleAttempt.stdout,
      peakMemoryMB: singleAttempt.peakMemoryMB,
      exitCode: singleAttempt.exitCode,
      promptRole: 'single-render'
    })
    throw new Error(singleAttempt.stderr || 'Manim render failed')
  }

  return {
    finalCode: initialCode,
    renderPeakMemoryMB: singleAttempt.peakMemoryMB
  }
}