import { type AnyRouter, useRouter, useSearch, type ValidateId } from '@tanstack/react-router'
import { useLayoutEffect, useMemo } from 'react'

import type { RegisteredRouter } from '#router/index'

type ConstrainLiteral<T, TConstraint, TDefault = TConstraint> = (T & TConstraint) | TDefault

type SearchSchema<
  TRouter extends AnyRouter,
  TFrom extends keyof TRouter['routesById'],
> = keyof TRouter['routesById'][TFrom]['types']['searchSchema']

type AnyKey<TRouter extends AnyRouter, TKey extends string> = ConstrainLiteral<
  TKey,
  string &
    {
      [K in keyof TRouter['routesById']]: SearchSchema<TRouter, K>
    }[keyof TRouter['routesById']]
>

type FromKey<
  TRouter extends AnyRouter,
  TFrom extends string,
  TKey extends string,
> = ConstrainLiteral<TKey, string & SearchSchema<TRouter, TFrom & keyof TRouter['routesById']>>

type Params<
  TRouter extends AnyRouter,
  TFrom extends string,
  TStrict extends boolean,
  TKey extends string,
  TValue,
> = TStrict extends false
  ? {
      from?: never
      strict: TStrict
      key: AnyKey<TRouter, TKey>
      initial?: TValue
    }
  : {
      from: ValidateId<TRouter, TFrom>
      strict?: TStrict
      key: FromKey<TRouter, TFrom, TKey>
      initial?: TValue
    }

type ValueFrom<
  TRouter extends AnyRouter,
  TFrom,
  TStrict extends boolean,
  TKey extends string,
> = TStrict extends false
  ?
      | undefined
      | {
          [K in keyof TRouter['routesById']]: TKey extends keyof TRouter['routesById'][K]['types']['searchSchema']
            ? TRouter['routesById'][K]['types']['searchSchema'][TKey]
            : never
        }[keyof TRouter['routesById']]
  : TRouter['routesById'][TFrom & keyof TRouter['routesById']]['types']['searchSchema'][TKey]

export function useSearchState<
  TRouter extends AnyRouter = RegisteredRouter,
  TFrom extends string = string,
  TStrict extends boolean = true,
  TKey extends string = string,
  TValue = ValueFrom<TRouter, TFrom, TStrict, TKey>,
>({
  from,
  strict,
  initial,
  key,
}: Params<TRouter, TFrom, TStrict, TKey, NoInfer<TValue>>): readonly [
  state: TValue,
  setState: (value: TValue | ((prev: TValue) => TValue), push_state?: boolean) => void,
] {
  // state
  const state = useSearch(
    useMemo(() => ({
      from,
      select: getter.bind(null, key),
      strict,
    }), [key, from, strict])
  ) as TValue

  // setState
  const router = useRouter()
  const setState = useMemo(() => (setter<TValue>).bind(null, router, key), [router, key])

  // initial value
  useLayoutEffect(() => {
    if (initial !== undefined && state === undefined) {
      setState(initial)
    }
  }, [])

  return [state, setState]
}

function getter(key: string, state: object) {
  // @ts-expect-error -- no need to strictly type this internal function
  return state[key]
}

const store = Symbol('search accumulator')
const push = Symbol('push accumulator')

function setter<T>(
  router: AnyRouter & { [store]?: object | null; [push]?: boolean | null },
  key: string,
  value: T | ((prev: T) => T),
  push_state?: boolean,
) {
  const prev = router[store] || router.state.location.search
  const next = {
    ...prev,
    [key]:
      typeof value === 'function'
        ? // @ts-expect-error -- no need to strictly type this internal function
          value(prev[key])
        : value,
  }

  if (next[key] === prev[key]) return

  router[store] = next
  const replace = !(router[push] ||= push_state)
  void router.navigate({
    hash: router.state.location.hash,
    replace,
    search: next,
    to: router.state.location.pathname,
  })
  /**
   * We temporarily store the search state in the router object
   * so that multiple calls to setState in the same tick will
   * accumulate changes in the same object instead of overwriting.
   *
   * We use a Symbol to avoid conflicts with other properties
   * that we don't own.
   *
   * After a tick, we nullify the property to avoid memory leaks,
   * and to ensure that the next setState call will start from
   * the current search state.
   */
  Promise.resolve().then(() => {
    router[store] = null
    router[push] = null
  })
}