import { z } from 'zod'
import type {
  ResponseContext,
  ValidatedContext,
  ValidationWarning,
  ZodValidationMode,
} from '../../types'
import { PipelineError } from '../../types'

export async function validateStage(
  ctx: ResponseContext,
  schema: z.ZodTypeAny,
  mode: ZodValidationMode = 'warn',
): Promise<ValidatedContext> {
  const result = schema.safeParse(ctx.response)

  if (result.success) {
    return { key: ctx.key, data: result.data as unknown, warnings: [] }
  }

  const warnings: ValidationWarning[] = result.error.issues.map((issue) => {
    const numericPath = issue.path.filter(
      (p): p is string | number =>
        typeof p === 'string' || typeof p === 'number',
    )
    return {
      path: numericPath.join('.'),
      message: issue.message,
      received: getByPath(ctx.response, numericPath),
    }
  })

  if (mode === 'error') {
    throw new PipelineError({
      phase: 'validate',
      cause: new Error(result.error.issues.map((i) => i.message).join('; ')),
      key: ctx.key,
    })
  }

  if (mode === 'passthrough') {
    // Surface warnings but keep the full raw data intact
    return { key: ctx.key, data: ctx.response, warnings }
  }

  // 'warn' (default): strip failing fields, return partial data
  const partialData = extractPartial(schema, ctx.response, result.error.issues)
  return { key: ctx.key, data: partialData, warnings }
}

function getByPath(obj: unknown, path: (string | number)[]): unknown {
  return path.reduce((cur: unknown, key) => {
    if (cur == null || typeof cur !== 'object') return undefined
    return (cur as Record<string | number, unknown>)[key]
  }, obj)
}

function extractPartial(
  schema: z.ZodTypeAny,
  input: unknown,
  issues: z.ZodIssue[],
): unknown {
  if (!(schema instanceof z.ZodObject)) return null
  if (typeof input !== 'object' || input === null) return null

  const failingKeys = new Set<string>()
  for (const issue of issues) {
    const first = issue.path[0]
    if (first != null && typeof first !== 'symbol')
      failingKeys.add(String(first))
  }

  const shape = schema.shape as Record<string, z.ZodTypeAny>
  const inputObj = input as Record<string, unknown>
  const partial: Record<string, unknown> = {}

  for (const key of Object.keys(shape)) {
    if (!failingKeys.has(key) && key in inputObj) {
      partial[key] = inputObj[key]
    }
  }

  return partial
}