import type { StateCreator, StoreApi, StoreMutatorIdentifier } from 'zustand'
import type { Transform, WidgetState } from './types'
import { applyTransforms } from './transforms'

type PipelineMiddleware = <
  T extends WidgetState,
  Mps extends [StoreMutatorIdentifier, unknown][] = [],
  Mcs extends [StoreMutatorIdentifier, unknown][] = [],
>(
  config: StateCreator<T, Mps, Mcs>,
) => StateCreator<T, Mps, Mcs>

type SetFn<T> = StoreApi<T>['setState']

// Per-store flag: true iff the current `state.error` was authored by the
// pipeline (a transform threw) rather than by the consumer (Provider's
// `error` prop or a direct setState). Lets the middleware clear stale
// pipeline errors on a subsequent clean run without clobbering errors a
// consumer set out-of-band.
const pipelineOwnedError = new WeakSet<StoreApi<WidgetState>>()

export const pipelineMiddleware: PipelineMiddleware = ((config) =>
  (set, get, api) => {
    const origSetState = api.setState as SetFn<WidgetState>
    const typedApi = api as StoreApi<WidgetState>

    const wrap = (delegate: SetFn<WidgetState>): SetFn<WidgetState> =>
      ((partial, replace) => {
        const current = (api.getState as () => WidgetState)()
        const incoming =
          typeof partial === 'function'
            ? (partial as (s: WidgetState) => Partial<WidgetState>)(current)
            : (partial as Partial<WidgetState>)

        const has = (key: keyof WidgetState) =>
          Object.prototype.hasOwnProperty.call(incoming, key)

        const rawDataChanged =
          has('rawData') && !Object.is(incoming.rawData, current.rawData)
        const dataXfChanged =
          has('dataTransforms') &&
          incoming.dataTransforms !== current.dataTransforms

        if (!rawDataChanged && !dataXfChanged) {
          // Consumer-authored writes that don't trigger the pipeline. If the
          // consumer is explicitly touching `error`, the next pipeline run is
          // no longer responsible for it — drop the pipeline-owned marker.
          if (has('error')) pipelineOwnedError.delete(typedApi)
          return delegate(incoming, replace as false)
        }

        const augmented: Partial<WidgetState> = { ...incoming }
        let firstError: unknown = undefined
        const onError = (_id: string, err: Error) => {
          if (firstError === undefined) firstError = err
        }

        const nextRaw = rawDataChanged ? incoming.rawData : current.rawData
        const nextXf: readonly Transform[] = dataXfChanged
          ? incoming.dataTransforms!
          : current.dataTransforms
        augmented.data = applyTransforms(nextRaw, nextXf, onError)

        // Reconcile `error` based on who authored the previous value:
        //   - pipeline throw → stamp the new error, mark pipeline-owned.
        //   - clean run AND previous error was pipeline-owned → clear it.
        //   - clean run AND previous error was consumer-set → leave it.
        // Consumer-authored `error` in `incoming` also wins over the marker
        // (and resets the marker via the early-return branch above).
        if (firstError !== undefined) {
          augmented.error = firstError
          pipelineOwnedError.add(typedApi)
        } else if (pipelineOwnedError.has(typedApi) && !has('error')) {
          augmented.error = undefined
          pipelineOwnedError.delete(typedApi)
        } else if (has('error')) {
          // Consumer set error in the same call that ran the pipeline — honor
          // their write and surrender ownership.
          pipelineOwnedError.delete(typedApi)
        }

        return delegate(augmented, replace as false)
      }) as SetFn<WidgetState>

    // Wrap both the public setState (used by store.setState() callers) and
    // the inner set passed into the state creator (used by actions inside config).
    api.setState = wrap(origSetState) as typeof api.setState
    const wrappedInnerSet = wrap(set as SetFn<WidgetState>)

    return config(wrappedInnerSet as typeof set, get, api)
  }) as PipelineMiddleware