import gettext
from contextlib import contextmanager
from gettext import GNUTranslations, NullTranslations
from importlib import import_module
from pathlib import Path
from typing import Any, Iterator

from jinja2 import pass_context
from jinja2.environment import Environment
from jinja2.loaders import FileSystemLoader
from jinja2.runtime import Context as JinjaContext

from jweb.content import Content, ContentField

_DEFAULT_LANGUAGE = "en"


class Context:
    def __init__(self, root_directory: Path) -> None:
        self.__root_directory = root_directory
        self.__output_directory = root_directory / "build"
        self.__content_directory = root_directory / "content"
        self.__environment = Environment(loader=FileSystemLoader(searchpath=root_directory / "src"))
        self.__translations: dict[str, GNUTranslations | NullTranslations] = {}
        self.__current_language: str | None = None

        self.add_filters(load=self.__load, write=self.__write, glob=self.__glob)
        self.add_globals(load=self.__load, write=self.__write, glob=self.__glob)

    @property
    def current_language(self) -> str | None:
        return self.__current_language

    def add_filters(self, **filters: Any) -> None:
        self.__environment.filters.update(**filters)

    def add_globals(self, **globals: Any) -> None:
        self.__environment.globals.update(**globals)

    def load_extensions(self, *extensions: str) -> None:
        for extension in extensions:
            module = import_module(extension)
            module.load_extension(self)

    def load_translations(self, domain: str, locale_dir: str | Path, *languages: str) -> None:
        locale_dir = str(locale_dir)
        self.__environment.add_extension("jinja2.ext.i18n")
        self.__translations[_DEFAULT_LANGUAGE] = NullTranslations()
        for language in languages:
            self.__translations[language] = gettext.translation(
                domain, localedir=str(locale_dir), languages=[language]
            )

    @contextmanager
    def set_language(self, language: str) -> Iterator[None]:
        translation = self.__translations[language]
        old_language = self.__current_language
        self.__environment.install_gettext_translations(translation, newstyle=True)  # type: ignore
        self.__current_language = language
        yield
        self.__current_language = old_language
        self.__environment.uninstall_gettext_translations(translation)  # type: ignore

    def render(self, source: str, output: str | Path, **context: Any) -> None:
        if self.__translations:
            for language in self.__translations:
                with self.set_language(language):
                    self.__render(source, self.__output_directory / language / output, **context)
        else:
            self.__render(source, self.__output_directory / output, **context)

    @pass_context
    def __load(self, context: JinjaContext, path: str | Path | ContentField) -> Content[bytes]:
        if isinstance(path, ContentField):
            path = path.as_path()

        current_language = self.current_language
        if current_language:
            localized_path = self.__content_directory / current_language / path
            if not localized_path.exists():
                localized_path = self.__content_directory / _DEFAULT_LANGUAGE / path
        else:
            localized_path = self.__content_directory / path

        with localized_path.open("rb") as content_file:
            return Content(localized_path, content_file.read(), current_language)

    @pass_context
    def __glob(
        self, context: JinjaContext, pattern: str, include_base_language: bool = False
    ) -> Iterator[Content[bytes]]:
        roots: list[Path] = []
        current_language = self.current_language
        if current_language is None:
            roots.append(self.__content_directory)
        else:
            if include_base_language:
                roots.append(self.__content_directory / _DEFAULT_LANGUAGE)
            roots.append(self.__content_directory / current_language)

        paths: set[Path] = set()
        for root in roots:
            paths = paths | set(it.relative_to(root) for it in root.glob(pattern))

        for path in paths:
            yield self.__load(context, path)

    def __write(self, content: Content[bytes]) -> Path:
        relative_path = content.path.relative_to(self.__content_directory)
        output_path = self.__output_directory / relative_path
        output_path.parent.mkdir(parents=True, exist_ok=True)
        with output_path.open("wb") as output_file:
            output_file.write(content.data)

        return Path(f"/{relative_path}")

    def __render(self, source: str, output_path: Path, **context: Any) -> None:
        output_path.parent.mkdir(exist_ok=True, parents=True)
        template = self.__environment.get_template(source)
        content = template.render(**context)

        with open(output_path, "w") as output_file:
            output_file.write(content)