import { QueryClient } from '@tanstack/react-query'
import type { z } from 'zod'
import { WorkflowEngine } from '../engine'
import type { FetchContext, SourceDescriptor } from '../types'

interface MockConfig {
  response?: unknown
  error?: Error
  handler?: (ctx: FetchContext) => Promise<unknown>
  delay?: number
}

interface TestEngineConfig {
  name: string
  mocks?: Record<string, MockConfig>
}

export interface TestWorkflowEngine
  extends WorkflowEngine<Record<string, SourceDescriptor<z.ZodTypeAny>>> {
  setMockResponse: (key: string, response: unknown) => void
  setMockError: (key: string, error: Error) => void
  setMockDelay: (key: string, ms: number) => void
  resetMocks: () => void
}

export function createTestEngine(
  config: TestEngineConfig = {},
): TestWorkflowEngine {
  const mocks = new Map<string, MockConfig>(Object.entries(config.mocks ?? {}))

  // staleTime:0, gcTime:0 so every trigger re-invokes the mock fn
  const queryClient = new QueryClient({
    defaultOptions: { queries: { retry: false, staleTime: 0, gcTime: 0 } },
  })

  // defaultThrottle:0 so test code can trigger the same key+params multiple
  // times without the 200ms default window swallowing subsequent calls.
  const engine = new WorkflowEngine({
    queryClient,
    defaultThrottle: 0,
    name: config.name,
  })

  function makeFetchFn(key: string): (ctx: FetchContext) => Promise<unknown> {
    return async (ctx: FetchContext): Promise<unknown> => {
      const mock = mocks.get(key)
      if (!mock) throw new Error(`No mock configured for key "${key}"`)
      if (mock.delay) await new Promise((r) => setTimeout(r, mock.delay))
      if (mock.error) throw mock.error
      if (mock.handler) return mock.handler(ctx)
      if (mock.response !== undefined) return mock.response
      throw new Error(`Mock for "${key}" has no response, error, or handler`)
    }
  }

  // Override register to auto-inject mock fetch and disable the default 200ms
  // throttle so tests can trigger the same key multiple times without the
  // throttle window swallowing subsequent calls. Descriptor-level throttle
  // still takes precedence if the test explicitly sets one.
  const originalRegister = engine.register.bind(engine)
  ;(engine as unknown as TestWorkflowEngine).register = <
    K extends string,
    S extends z.ZodTypeAny,
  >(
    descriptor: SourceDescriptor<S> & { key: K },
  ) =>
    originalRegister({
      throttle: 0,
      ...descriptor,
      fetch: makeFetchFn(descriptor.key),
    })

  const testEngine = engine as unknown as TestWorkflowEngine

  testEngine.setMockResponse = (key, response) => {
    mocks.set(key, { ...(mocks.get(key) ?? {}), response, error: undefined })
    queryClient.removeQueries({ queryKey: [key] })
  }

  testEngine.setMockError = (key, error) => {
    mocks.set(key, { ...(mocks.get(key) ?? {}), error, response: undefined })
    queryClient.removeQueries({ queryKey: [key] })
  }

  testEngine.setMockDelay = (key, ms) => {
    mocks.set(key, { ...(mocks.get(key) ?? {}), delay: ms })
  }

  testEngine.resetMocks = () => {
    mocks.clear()
    queryClient.clear()
  }

  return testEngine
}