import { createLogger } from '@sim/logger'
import { createTimeoutAbortController, getTimeoutErrorMessage } from '@/lib/core/execution-limits'
import {
  extractBlockIdFromOutputId,
  extractPathFromOutputId,
  traverseObjectPath,
} from '@/lib/core/utils/response-format'
import { encodeSSE } from '@/lib/core/utils/sse'
import { compactExecutionPayload } from '@/lib/execution/payloads/serializer'
import { buildTraceSpans } from '@/lib/logs/execution/trace-spans/trace-spans'
import { processStreamingBlockLogs } from '@/lib/tokenization'
import {
  cleanupExecutionBase64Cache,
  hydrateUserFilesWithBase64,
} from '@/lib/uploads/utils/user-file-base64.server'
import type { BlockLog, ExecutionResult, StreamingExecution } from '@/executor/types'

/**
 * Extended streaming execution type that includes blockId on the execution.
 * The runtime passes blockId but the base StreamingExecution type doesn't declare it.
 */
interface StreamingExecutionWithBlockId extends Omit<StreamingExecution, 'execution'> {
  execution?: StreamingExecution['execution'] & { blockId?: string }
}

const logger = createLogger('WorkflowStreaming')

const DANGEROUS_KEYS = ['__proto__', 'constructor', 'prototype']

interface StreamingConfig {
  selectedOutputs?: string[]
  isSecureMode?: boolean
  workflowTriggerType?: 'api' | 'chat'
  includeFileBase64?: boolean
  base64MaxBytes?: number
  timeoutMs?: number
}

export type StreamingExecutorFn = (callbacks: {
  onStream: (streamingExec: StreamingExecution) => Promise<void>
  onBlockComplete: (blockId: string, output: unknown) => Promise<void>
  abortSignal: AbortSignal
}) => Promise<ExecutionResult>

export interface StreamingResponseOptions {
  requestId: string
  streamConfig: StreamingConfig
  executionId?: string
  largeValueExecutionIds?: string[]
  allowLargeValueWorkflowScope?: boolean
  workspaceId?: string
  workflowId?: string
  userId?: string
  executeFn: StreamingExecutorFn
}

interface StreamingState {
  streamedChunks: Map<string, string[]>
  processedOutputs: Set<string>
  streamCompletionTimes: Map<string, number>
  completedBlockIds: Set<string>
}

function resolveStreamedContent(state: StreamingState): Map<string, string> {
  const result = new Map<string, string>()
  for (const [blockId, chunks] of state.streamedChunks) {
    result.set(blockId, chunks.join(''))
  }
  return result
}

function extractOutputValue(output: unknown, path: string): unknown {
  return traverseObjectPath(output, path)
}

function isDangerousKey(key: string): boolean {
  return DANGEROUS_KEYS.includes(key)
}

async function buildMinimalResult(
  result: ExecutionResult,
  selectedOutputs: string[] | undefined,
  streamedContent: Map<string, string>,
  completedBlockIds: Set<string>,
  requestId: string,
  includeFileBase64: boolean,
  base64MaxBytes: number | undefined,
  executionId?: string,
  context: Pick<StreamingResponseOptions, 'workspaceId' | 'workflowId' | 'userId'> = {}
): Promise<{ success: boolean; error?: string; output: Record<string, unknown> }> {
  const durableContext = {
    workspaceId: context.workspaceId,
    workflowId: context.workflowId,
    executionId,
    userId: context.userId,
    requireDurable: Boolean(context.workspaceId && context.workflowId && executionId),
  }

  const minimalResult = {
    success: result.success,
    error: result.error,
    output: {} as Record<string, unknown>,
  }

  if (result.status === 'paused') {
    minimalResult.output = result.output || {}
    return compactExecutionPayload(minimalResult, {
      ...durableContext,
      preserveUserFileBase64: includeFileBase64,
      preserveRoot: true,
    })
  }

  if (!selectedOutputs?.length) {
    minimalResult.output = result.output || {}
    return compactExecutionPayload(minimalResult, {
      ...durableContext,
      preserveUserFileBase64: includeFileBase64,
      preserveRoot: true,
    })
  }

  if (!result.output || !result.logs) {
    return minimalResult
  }

  for (const outputId of selectedOutputs) {
    const blockId = extractBlockIdFromOutputId(outputId)

    if (streamedContent.has(blockId)) {
      continue
    }

    if (!completedBlockIds.has(blockId)) {
      continue
    }

    if (isDangerousKey(blockId)) {
      logger.warn(`[${requestId}] Blocked dangerous blockId: ${blockId}`)
      continue
    }

    const path = extractPathFromOutputId(outputId, blockId)
    if (isDangerousKey(path)) {
      logger.warn(`[${requestId}] Blocked dangerous path: ${path}`)
      continue
    }

    const blockLog = result.logs.find((log: BlockLog) => log.blockId === blockId)
    if (!blockLog?.output) {
      continue
    }

    const value = extractOutputValue(blockLog.output, path)
    if (value === undefined) {
      continue
    }

    if (!minimalResult.output[blockId]) {
      minimalResult.output[blockId] = Object.create(null) as Record<string, unknown>
    }
    ;(minimalResult.output[blockId] as Record<string, unknown>)[path] = value
  }

  return compactExecutionPayload(minimalResult, {
    ...durableContext,
    preserveUserFileBase64: includeFileBase64,
    preserveRoot: true,
  })
}

function updateLogsWithStreamedContent(
  logs: BlockLog[],
  streamedContent: Map<string, string>,
  streamCompletionTimes: Map<string, number>
): BlockLog[] {
  return logs.map((log: BlockLog) => {
    if (!streamedContent.has(log.blockId)) {
      return log
    }

    const content = streamedContent.get(log.blockId)
    const updatedLog = { ...log }

    if (streamCompletionTimes.has(log.blockId)) {
      const completionTime = streamCompletionTimes.get(log.blockId)!
      const startTime = new Date(log.startedAt).getTime()
      updatedLog.endedAt = new Date(completionTime).toISOString()
      updatedLog.durationMs = completionTime - startTime
    }

    if (log.output && content) {
      updatedLog.output = { ...log.output, content }
    }

    return updatedLog
  })
}

async function completeLoggingSession(result: ExecutionResult): Promise<void> {
  if (!result._streamingMetadata?.loggingSession) {
    return
  }

  const { traceSpans, totalDuration } = buildTraceSpans(result)

  await result._streamingMetadata.loggingSession.safeComplete({
    endedAt: new Date().toISOString(),
    totalDurationMs: totalDuration || 0,
    finalOutput: result.output || {},
    traceSpans: (traceSpans || []) as any,
    workflowInput: result._streamingMetadata.processedInput,
  })

  result._streamingMetadata = undefined
}

export async function createStreamingResponse(
  options: StreamingResponseOptions
): Promise<ReadableStream> {
  const { requestId, streamConfig, executionId, executeFn } = options
  const durableContext = {
    workspaceId: options.workspaceId,
    workflowId: options.workflowId,
    executionId,
    userId: options.userId,
    requireDurable: Boolean(options.workspaceId && options.workflowId && executionId),
  }
  const timeoutController = createTimeoutAbortController(streamConfig.timeoutMs)

  return new ReadableStream({
    async start(controller) {
      const state: StreamingState = {
        streamedChunks: new Map(),
        processedOutputs: new Set(),
        streamCompletionTimes: new Map(),
        completedBlockIds: new Set(),
      }

      const sendChunk = (blockId: string, content: string) => {
        const separator = state.processedOutputs.size > 0 ? '\n\n' : ''
        controller.enqueue(encodeSSE({ blockId, chunk: separator + content }))
        state.processedOutputs.add(blockId)
      }

      /**
       * Callback for handling streaming execution events.
       */
      const onStreamCallback = async (streamingExec: StreamingExecutionWithBlockId) => {
        const blockId = streamingExec.execution?.blockId
        if (!blockId) {
          logger.warn(`[${requestId}] Streaming execution missing blockId`)
          return
        }

        const reader = streamingExec.stream.getReader()
        const decoder = new TextDecoder()
        let isFirstChunk = true

        try {
          while (true) {
            const { done, value } = await reader.read()
            if (done) {
              state.streamCompletionTimes.set(blockId, Date.now())
              break
            }

            const textChunk = decoder.decode(value, { stream: true })
            if (!state.streamedChunks.has(blockId)) {
              state.streamedChunks.set(blockId, [])
            }
            state.streamedChunks.get(blockId)!.push(textChunk)

            if (isFirstChunk) {
              sendChunk(blockId, textChunk)
              isFirstChunk = false
            } else {
              controller.enqueue(encodeSSE({ blockId, chunk: textChunk }))
            }
          }
        } catch (error) {
          logger.error(`[${requestId}] Error reading stream for block ${blockId}:`, error)
          controller.enqueue(
            encodeSSE({
              event: 'stream_error',
              blockId,
              error: error instanceof Error ? error.message : 'Stream reading error',
            })
          )
        }
      }

      const includeFileBase64 = streamConfig.includeFileBase64 ?? true
      const base64MaxBytes = streamConfig.base64MaxBytes

      const onBlockCompleteCallback = async (blockId: string, output: unknown) => {
        state.completedBlockIds.add(blockId)

        if (!streamConfig.selectedOutputs?.length) {
          return
        }

        if (state.streamedChunks.has(blockId)) {
          return
        }

        const matchingOutputs = streamConfig.selectedOutputs.filter(
          (outputId) => extractBlockIdFromOutputId(outputId) === blockId
        )

        for (const outputId of matchingOutputs) {
          const path = extractPathFromOutputId(outputId, blockId)
          const outputValue = extractOutputValue(output, path)

          if (outputValue !== undefined) {
            const hydratedOutput = includeFileBase64
              ? await hydrateUserFilesWithBase64(outputValue, {
                  requestId,
                  workspaceId: options.workspaceId,
                  workflowId: options.workflowId,
                  executionId,
                  largeValueExecutionIds: options.largeValueExecutionIds,
                  allowLargeValueWorkflowScope: options.allowLargeValueWorkflowScope,
                  userId: options.userId,
                  maxBytes: base64MaxBytes,
                })
              : outputValue
            const compactHydratedOutput = await compactExecutionPayload(hydratedOutput, {
              ...durableContext,
              preserveUserFileBase64: includeFileBase64,
            })
            const formattedOutput =
              typeof compactHydratedOutput === 'string'
                ? compactHydratedOutput
                : JSON.stringify(compactHydratedOutput, null, 2)
            sendChunk(blockId, formattedOutput)
          }
        }
      }

      try {
        const result = await executeFn({
          onStream: onStreamCallback,
          onBlockComplete: onBlockCompleteCallback,
          abortSignal: timeoutController.signal,
        })

        const streamedContent =
          state.streamedChunks.size > 0 ? resolveStreamedContent(state) : new Map<string, string>()

        if (result.logs && streamedContent.size > 0) {
          result.logs = updateLogsWithStreamedContent(
            result.logs,
            streamedContent,
            state.streamCompletionTimes
          )
          processStreamingBlockLogs(result.logs, streamedContent)
        }

        if (
          result.status === 'cancelled' &&
          timeoutController.isTimedOut() &&
          timeoutController.timeoutMs
        ) {
          const timeoutErrorMessage = getTimeoutErrorMessage(null, timeoutController.timeoutMs)
          logger.info(`[${requestId}] Streaming execution timed out`, {
            timeoutMs: timeoutController.timeoutMs,
          })
          if (result._streamingMetadata?.loggingSession) {
            await result._streamingMetadata.loggingSession.markAsFailed(timeoutErrorMessage)
          }
          controller.enqueue(encodeSSE({ event: 'error', error: timeoutErrorMessage }))
        } else {
          await completeLoggingSession(result)

          const minimalResult = await buildMinimalResult(
            result,
            streamConfig.selectedOutputs,
            streamedContent,
            state.completedBlockIds,
            requestId,
            streamConfig.includeFileBase64 ?? true,
            streamConfig.base64MaxBytes,
            executionId,
            {
              workspaceId: options.workspaceId,
              workflowId: options.workflowId,
              userId: options.userId,
            }
          )

          controller.enqueue(
            encodeSSE({
              event: 'final',
              data: {
                ...minimalResult,
                ...(result.status === 'paused' && { status: 'paused' }),
              },
            })
          )
        }

        controller.enqueue(encodeSSE('[DONE]'))

        if (executionId) {
          await cleanupExecutionBase64Cache(executionId)
        }

        controller.close()
      } catch (error: any) {
        logger.error(`[${requestId}] Stream error:`, error)
        controller.enqueue(
          encodeSSE({ event: 'error', error: error.message || 'Stream processing error' })
        )

        if (executionId) {
          await cleanupExecutionBase64Cache(executionId)
        }

        controller.close()
      } finally {
        timeoutController.cleanup()
      }
    },
    async cancel(reason) {
      logger.info(`[${requestId}] Streaming response cancelled`, { reason })
      timeoutController.abort()
      timeoutController.cleanup()
      if (executionId) {
        try {
          await cleanupExecutionBase64Cache(executionId)
        } catch (error) {
          logger.error(`[${requestId}] Failed to cleanup base64 cache`, { error })
        }
      }
    },
  })
}
