import React, { createContext, useCallback, useContext, useEffect, useMemo, useState } from 'react';
import { z } from 'zod';
import { en } from '@/locales/en.ts';
import { ru } from '@/locales/ru.ts';
import { tt } from '@/locales/tt.ts';

type Language = 'en' | 'ru' | 'tt';

// Translation value can be string, object, or array
const TranslationValueSchema: z.ZodTypeAny = z.lazy(() =>
  z.union([z.string(), z.record(TranslationValueSchema), z.array(TranslationValueSchema)])
);

// Translation object schema
// eslint-disable-next-line @typescript-eslint/no-unused-vars
const TranslationSchema = z.record(TranslationValueSchema);

// Replacements schema for interpolation
const ReplacementsSchema = z
  .object({
    count: z.number().optional(),
  })
  .catchall(z.union([z.string(), z.number()]));

interface I18nContextType {
  lang: Language;
  setLang: (lang: Language) => void;
  t: (key: string, replacements?: z.infer<typeof ReplacementsSchema>) => string;
  t_raw: (key: string) => z.infer<typeof TranslationValueSchema>;
}

const translations = { en, ru, tt };
const LANG_STORAGE_KEY = 'tugan-yak-lang';

const I18nContext = createContext<I18nContextType | undefined>(undefined);

// Helper to resolve nested keys like 'a.b.c'
const resolvePath = (
  obj: z.infer<typeof TranslationSchema>,
  path: string
): z.infer<typeof TranslationValueSchema> => {
  return path.split('.').reduce((acc: z.infer<typeof TranslationValueSchema>, part: string) => {
    if (acc && typeof acc === 'object' && !Array.isArray(acc)) {
      return (acc as Record<string, z.infer<typeof TranslationValueSchema>>)[part];
    }
    return undefined;
  }, obj);
};

export const I18nProvider = ({ children }: { children?: React.ReactNode }) => {
  const [lang, setLang] = useState<Language>(() => {
    try {
      const storedLang = window.localStorage.getItem(LANG_STORAGE_KEY);
      if (storedLang === 'en' || storedLang === 'ru' || storedLang === 'tt') {
        return storedLang;
      }
      const browserLang = navigator.language.split(/[-_]/)[0];
      if (browserLang === 'ru') return 'ru';
      if (browserLang === 'tt') return 'tt';
    } catch (e) {
      console.error('Could not read language preference.', e);
    }
    return 'en'; // Default to English
  });

  useEffect(() => {
    document.documentElement.lang = lang;
    try {
      window.localStorage.setItem(LANG_STORAGE_KEY, lang);
    } catch (e) {
      console.error('Could not save language to localStorage', e);
    }
  }, [lang]);

  const t_raw = useCallback(
    (key: string): z.infer<typeof TranslationValueSchema> => {
      let translation = resolvePath(translations[lang], key);
      if (translation === undefined) {
        // Fallback to English if translation is missing in the current language
        translation = resolvePath(translations.en, key);
      }
      return translation;
    },
    [lang]
  );

  // Memoize plural rules to avoid recreating on every call
  const pluralRules = useMemo(() => new Intl.PluralRules(lang), [lang]);

  const t = useCallback(
    (key: string, replacements?: z.infer<typeof ReplacementsSchema>): string => {
      let translation: z.infer<typeof TranslationValueSchema>;
      let usedKey = key;

      if (replacements?.count !== undefined) {
        const pluralCategory = pluralRules.select(replacements.count);

        const pluralKey = `${key}.${pluralCategory}`;
        const pluralTranslation = t_raw(pluralKey);

        if (typeof pluralTranslation === 'string') {
          translation = pluralTranslation;
          usedKey = pluralKey;
        } else {
          // Fallback to the '.other' key which is standard for i18next-style pluralization
          const otherKey = `${key}.other`;
          const otherTranslation = t_raw(otherKey);
          if (typeof otherTranslation === 'string') {
            translation = otherTranslation;
            usedKey = otherKey;
          } else {
            // Final fallback to non-plural version
            translation = t_raw(key);
          }
        }
      } else {
        translation = t_raw(key);
      }

      if (typeof translation !== 'string') {
        console.warn(
          `Translation for key "${usedKey}" is not a string. Found: ${typeof translation}. Falling back to key.`
        );
        return key;
      }

      if (replacements) {
        // Replace placeholders - create regex only once per placeholder
        let result = translation;
        Object.keys(replacements).forEach((placeholder) => {
          const regex = new RegExp(`{{${placeholder}}}`, 'g');
          result = result.replace(regex, String(replacements[placeholder]));
        });
        return result;
      }

      return translation;
    },
    [t_raw, pluralRules]
  );

  return (
    <I18nContext.Provider value={{ lang, setLang, t, t_raw }}>{children}</I18nContext.Provider>
  );
};

export const useTranslation = (): I18nContextType => {
  const context = useContext(I18nContext);
  if (!context) {
    throw new Error('useTranslation must be used within an I18nProvider');
  }
  return context;
};