Momento/memento-note/lib/i18n/LanguageProvider.tsx

'use client'

import { createContext, useContext, useEffect, useState, useCallback, useRef, useMemo } from 'react'
import type { ReactNode } from 'react'
import { SupportedLanguage, loadTranslations, getTranslationValue, Translations } from './load-translations'

// Static imports for SSR-safe initial translations (prevents hydration mismatch)
import enTranslations from '@/locales/en.json'

type LanguageContextType = {
  language: SupportedLanguage
  setLanguage: (lang: SupportedLanguage) => void
  t: (key: string, params?: Record<string, string | number>) => string
  translations: Translations
}

const LanguageContext = createContext<LanguageContextType | undefined>(undefined)

const RTL_LANGUAGES: SupportedLanguage[] = ['ar', 'fa']
const SUPPORTED_LANGS: SupportedLanguage[] = ['en', 'fr', 'es', 'de', 'fa', 'it', 'pt', 'ru', 'zh', 'ja', 'ko', 'ar', 'hi', 'nl', 'pl']

function updateDocumentDirection(lang: SupportedLanguage) {
  document.documentElement.lang = lang
  document.documentElement.dir = RTL_LANGUAGES.includes(lang) ? 'rtl' : 'ltr'
}



export function LanguageProvider({ children, initialLanguage = 'en', initialTranslations }: {
  children: ReactNode
  initialLanguage?: SupportedLanguage
  initialTranslations?: Translations
}) {
  // Use server-provided initialLanguage and translations for the first render 
  // to ensure client hydration matches server HTML exactly.
  const [language, setLanguageState] = useState<SupportedLanguage>(initialLanguage)

  // Start with server-provided translations or fallback to English ONLY if it's the requested language
  const [translations, setTranslations] = useState<Translations>(() => {
    if (initialTranslations) return initialTranslations
    return enTranslations as unknown as Translations
  })

  const cacheRef = useRef<Map<SupportedLanguage, Translations>>(new Map())
  
  // Initialize cache with initial translations to prevent redundant loading
  if (initialTranslations && !cacheRef.current.has(initialLanguage)) {
    cacheRef.current.set(initialLanguage, initialTranslations)
  }

  const isFirstRender = useRef(true)

  // Load saved preference from localStorage AFTER hydration
  useEffect(() => {
    const saved = localStorage.getItem('user-language') as SupportedLanguage
    if (saved && SUPPORTED_LANGS.includes(saved) && saved !== initialLanguage) {
      setLanguageState(saved)
    }
  }, [initialLanguage])

  // Always reload locale JSON on the client so new i18n keys are picked up
  // (SSR initialTranslations may be stale until the next deploy).
  useEffect(() => {
    let cancelled = false

    const loadLang = async () => {
      const loaded = await loadTranslations(language)
      if (cancelled) return
      cacheRef.current.set(language, loaded)
      setTranslations(loaded)
      if (!isFirstRender.current) updateDocumentDirection(language)
      isFirstRender.current = false
    }

    loadLang()
    return () => {
      cancelled = true
    }
  }, [language])

  const setLanguage = useCallback((lang: SupportedLanguage) => {
    setLanguageState(lang)
    localStorage.setItem('user-language', lang)
    document.cookie = `user-language=${lang};path=/;max-age=${60 * 60 * 24 * 365};samesite=lax`
    updateDocumentDirection(lang)
  }, [])

  const t = useCallback((key: string, params?: Record<string, string | number>) => {
    if (!translations) return key

    let value: any = getTranslationValue(translations, key)
    if (value === key && language !== 'en') {
      value = getTranslationValue(enTranslations as unknown as Translations, key)
    }

    if (params && typeof value === 'string') {
      Object.entries(params).forEach(([param, paramValue]) => {
        value = value.replace(`{${param}}`, String(paramValue))
      })
    }

    return typeof value === 'string' ? value : key
  }, [translations, language])

  const value = useMemo(() => ({ language, setLanguage, t, translations }), [language, setLanguage, t, translations])

  return (
    <LanguageContext.Provider value={value}>
      {children}
    </LanguageContext.Provider>
  )
}

export function useLanguage() {
  const context = useContext(LanguageContext)
  if (context === undefined) {
    throw new Error('useLanguage must be used within a LanguageProvider')
  }
  return context
}