1import { createContext, useContext, useEffect, useState } from "react"
 2
 3export type Theme = "dark" | "light" | "system"
 4
 5type ThemeProviderProps = {
 6    children: React.ReactNode
 7    defaultTheme?: Theme
 8    storageKey?: string
 9}
10
11type ThemeProviderState = {
12    theme: Theme
13    setTheme: (theme: Theme) => void
14}
15
16const initialState: ThemeProviderState = {
17    theme: "system",
18    setTheme: () => null,
19}
20
21const ThemeProviderContext = createContext<ThemeProviderState>(initialState)
22
23export function ThemeProvider({
24    children,
25    defaultTheme = "system",
26    storageKey = "vite-ui-theme",
27    ...props
28}: ThemeProviderProps) {
29    const [theme, setTheme] = useState<Theme>(
30        () => (localStorage.getItem(storageKey) as Theme) || defaultTheme
31    )
32
33    useEffect(() => {
34        const root = window.document.documentElement
35
36        root.classList.remove("light", "dark")
37
38        if (theme === "system") {
39            const systemTheme = window.matchMedia("(prefers-color-scheme: dark)")
40                .matches
41                ? "dark"
42                : "light"
43
44            root.classList.add(systemTheme)
45            return
46        }
47
48        root.classList.add(theme)
49    }, [theme])
50
51    const value = {
52        theme,
53        setTheme: (theme: Theme) => {
54            localStorage.setItem(storageKey, theme)
55            setTheme(theme)
56        },
57    }
58
59    return (
60        <ThemeProviderContext.Provider {...props} value={value}>
61            {children}
62        </ThemeProviderContext.Provider>
63    )
64}
65
66export const useTheme = () => {
67    const context = useContext(ThemeProviderContext)
68
69    if (context === undefined)
70        throw new Error("useTheme must be used within a ThemeProvider")
71
72    return context
73}