Tanstack Start React Themes

import { createServerFn } from "@tanstack/react-start";
import { getCookie, setCookie } from "@tanstack/react-start/server";
import { schema, ThemedefinedByMap, type ThemeIdentifier } from "./schema";

export const storageKey = "ui-theme";

export const getThemeServerFn = createServerFn().handler(async () => {
    const themeCookie = getCookie(storageKey);

    if (!themeCookie) {
        return null;
    }

    // Type guard to ensure themeCookie is a valid ThemeIdentifier
    if (!(themeCookie in ThemedefinedByMap)) {
        throw new Error("Invalid theme exists");
    }

    const { success, data } = schema.safeParse(ThemedefinedByMap[themeCookie as ThemeIdentifier]);
    if (!success) {
        throw new Error("Invalid theme exists");
    }

    return data;
});

export const setThemeServerFn = createServerFn({ method: "POST" })
    .validator((data: { identifier: keyof typeof ThemedefinedByMap }) => {
        const { success } = schema.safeParse(ThemedefinedByMap[data.identifier]);
        if (!success) {
            throw new Error("Invalid theme provided");
        }
        return data.identifier;
    })
    .handler(async ({ data }) => {
        setCookie(storageKey, data, {
            maxAge: 60 * 60 * 24 * 30 * 12, // 12 months
            httpOnly: true,
            secure: process.env.NODE_ENV === "production",
        });
    });

import { z } from "zod";

export const ThemedefinedByMap = Object.freeze({
    ud: { definedBy: 'user', theme: 'dark' },
    sd: { definedBy: 'system', theme: 'dark' },
    ul: { definedBy: 'user', theme: 'light' },
    sl: { definedBy: 'system', theme: 'light' },
});

export type ThemeIdentifier = keyof typeof ThemedefinedByMap;

// NOTE: Default to sd
export const schema = z.object({
    theme: z.enum(["dark", "light"]).optional().default('dark'),
    definedBy: z.enum(['user', 'system']).optional().default('system'),
});

export type ThemeType = z.infer<typeof schema>;

"use client";

import { createContext, useContext, useEffect } from "react";
import { setThemeServerFn } from "./action";
import { ThemeType, ThemeIdentifier } from "./schema";
import { useMediaQuery } from "../../hooks/useMediaQuery";
import { useRouter } from "@tanstack/react-router";

const COLOR_SCHEME_QUERY = '(prefers-color-scheme: dark)';

type SetThemeArg = ThemeType['theme'] | Extract<ThemeType['definedBy'], 'system'>

type ThemeProviderProps = {
    children: React.ReactNode
    value: ThemeType | null
    defaultValue?: ThemeType
}

type ThemeProviderState = {
    theme: ThemeType['theme']
    definedBy: ThemeType['definedBy']
    setTheme: (themeIdentifier: SetThemeArg) => void
}

const initialState: ThemeProviderState = {
    theme: 'dark',
    definedBy: 'system',
    setTheme: () => null,
}

const ThemeProviderContext = createContext<ThemeProviderState>(initialState)

export function ThemeProvider({
    children,
    value,
    defaultValue = { definedBy: 'system', theme: 'dark' },
    ...props
}: ThemeProviderProps) {
    const router = useRouter();
    const matches = useMediaQuery(COLOR_SCHEME_QUERY);
    const matchedMode = matches ? 'dark' : 'light'

    const themeCookieSetter = async (identifier: ThemeIdentifier) => {
        await setThemeServerFn({ data: { identifier } })
        router.invalidate()
    }

    useEffect(() => {
        const root = window.document.documentElement
        root.classList.remove("light", "dark")

        if (!value) {
            root.classList.add(matchedMode)
            themeCookieSetter(`s${matchedMode[0]}` as ThemeIdentifier)
            return
        }

        if (value?.definedBy === 'system' && matchedMode !== value.theme) {
            root.classList.add(matchedMode)

            const identifier = `s${matchedMode[0]}` as ThemeIdentifier
            themeCookieSetter(identifier)
            return
        }

        root.classList.add(value.theme)
    }, [value, matches])

    const _value = {
        theme: value?.theme || defaultValue.theme,
        definedBy: value?.definedBy || defaultValue.definedBy,
        setTheme: (args: SetThemeArg) => {
            let identifier: ThemeIdentifier
            if (args === 'system') {
                const systemTheme = matches ? 'dark' : 'light'
                identifier = `${args[0]}${systemTheme[0]}` as ThemeIdentifier
            } else {
                identifier = `u${args[0]}` as ThemeIdentifier
            }
            themeCookieSetter(identifier)
        },
    }

    return (
        <ThemeProviderContext.Provider {...props} value={_value}>
            {children}
        </ThemeProviderContext.Provider>
    )
}

export const useTheme = () => {
    const context = useContext(ThemeProviderContext)

    if (context === undefined)
        throw new Error("useTheme must be used within a ThemeProvider")

    return context
}

export * from './action'
export * from './schema'
export * from './provider'
export * from '../../hooks/useMediaQuery'
export * from '../../hooks/useIsomorphicLayoutEffect'