import { useState, useRef, useMemo, useEffect, useCallback } from "react"
import { produce, Draft } from "immer"

/**
 * Context provided to the action factory function.
 * Gives actions access to read and write state.
 */
export interface ActionContext<S> {
  /** Get the current state. Always returns the latest value, even mid-async-action. */
  getState(): S
  /** Replace the state entirely. Triggers a React re-render. */
  setState(state: S): void
  /** Update state using an immer recipe — mutate the draft freely. Triggers a React re-render. */
  update(recipe: (draft: Draft<S>) => void): void
}

/**
 * Maps each action function to a queued version that returns a Promise.
 * Async actions preserve their return type. Sync actions get wrapped in Promise.
 */
export type QueuedActions<A> = {
  [K in keyof A]: A[K] extends (...args: infer P) => Promise<infer R>
    ? (...args: P) => Promise<R>
    : A[K] extends (...args: infer P) => infer R
      ? (...args: P) => Promise<R>
      : never
}

/** Entry in the action queue */
interface QueueEntry {
  execute: () => Promise<any>
  resolve: (value: any) => void
  reject: (error: any) => void
}

/**
 * A typesafe React hook for local state management with:
 * - Immer-powered immutable updates
 * - Async action support
 * - Serialized (one-at-a-time) action execution via a FIFO queue
 * - Fully typesafe actions — no action constants needed
 *
 * Actions are defined as plain functions in a factory closure. The factory
 * receives an {@link ActionContext} for reading and writing state.
 *
 * @param initialState - initial state value, or a lazy initializer function
 * @param actionFactory - function that receives an ActionContext and returns
 *   a record of action functions
 * @returns a tuple of `[state, actions]` where actions are queued wrappers
 *   of the originals
 *
 * @example
 * ```ts
 * const [state, actions] = useStateMachine(
 *   { count: 0, loading: false },
 *   (ctx) => ({
 *     increment() {
 *       ctx.update(draft => { draft.count++ })
 *     },
 *     async loadData(url: string) {
 *       ctx.update(draft => { draft.loading = true })
 *       const data = await fetch(url).then(r => r.json())
 *       ctx.update(draft => {
 *         draft.count = data.count
 *         draft.loading = false
 *       })
 *       return data
 *     }
 *   })
 * )
 * ```
 */
export function useStateMachine<S, A extends Record<string, (...args: any[]) => any>>(
  initialState: S | (() => S),
  actionFactory: (ctx: ActionContext<S>) => A
): [S, QueuedActions<A>] {
  const [state, setReactState] = useState<S>(initialState)

  // Ref always holds the latest state so getState() works inside async actions
  const stateRef = useRef(state)
  stateRef.current = state

  const unmountedRef = useRef(false)
  const queueRef = useRef<QueueEntry[]>([])
  const runningRef = useRef(false)

  // Stable function to update React state, respecting unmount
  const commitState = useCallback((newState: S) => {
    stateRef.current = newState
    if (!unmountedRef.current) {
      setReactState(newState)
    }
  }, [])

  // Process the action queue serially
  const processQueue = useCallback(async () => {
    if (runningRef.current) return
    runningRef.current = true

    while (queueRef.current.length > 0) {
      if (unmountedRef.current) {
        // Reject everything remaining
        for (const entry of queueRef.current) {
          entry.reject(new Error("Component unmounted"))
        }
        queueRef.current = []
        break
      }

      const entry = queueRef.current.shift()!
      try {
        const result = await entry.execute()
        entry.resolve(result)
      } catch (err) {
        entry.reject(err)
        // Stop the queue on error — reject all remaining entries
        for (const remaining of queueRef.current) {
          remaining.reject(new Error("Action queue stopped due to a previous error"))
        }
        queueRef.current = []
        break
      }
    }

    runningRef.current = false
  }, [])

  // Build the ActionContext (stable, never changes)
  const ctx: ActionContext<S> = useMemo(() => ({
    getState: () => stateRef.current,
    setState: (newState: S) => commitState(newState),
    update: (recipe: (draft: Draft<S>) => void) => {
      commitState(produce(stateRef.current, recipe))
    }
  }), [commitState])

  // Create the raw actions once from the factory
  const rawActions = useMemo(() => actionFactory(ctx), [actionFactory, ctx])

  // Wrap each action so it goes through the queue
  const queuedActions = useMemo(() => {
    const wrapped: Record<string, (...args: any[]) => Promise<any>> = {}
    for (const key of Object.keys(rawActions)) {
      wrapped[key] = (...args: any[]) => {
        return new Promise((resolve, reject) => {
          queueRef.current.push({
            execute: () => Promise.resolve(rawActions[key](...args)),
            resolve,
            reject
          })
          processQueue()
        })
      }
    }
    return wrapped as QueuedActions<A>
  }, [rawActions, processQueue])

  // Mark unmounted on cleanup
  useEffect(() => {
    return () => {
      unmountedRef.current = true
    }
  }, [])

  return [state, queuedActions]
}