"""Jinja bytecode cache.

Monkey-patches get_template() to cache compiled bytecode on disk
using marshal. On cache hit, skips Jinja's lex/parse/codegen/compile
pipeline entirely.
"""

from __future__ import annotations

import marshal
import time
from importlib.metadata import version
from pathlib import Path

import jinja2
from hash_utils import fnv64

from dbt_run_cache.config import DBT_RUN_CACHE_PATH

_CACHE_MAX_AGE_SECONDS = 14 * 86400  # 2 weeks
_VERSION_SALT = "__".join(
    (
        "dbt" + version("dbt-core").replace(".", "_"),
        "j" + version("jinja2").replace(".", "_"),
    )
)
_CACHE_DIR = DBT_RUN_CACHE_PATH / "jinja_cache"


def _jinja_hash(source: str) -> str:
    """FNV-1a 64-bit hash of template source with version salt."""
    return f"{_VERSION_SALT}_{fnv64(source.encode())}"


def _clean_stale_files(cache_dir: Path) -> None:
    """Delete cache files older than 2 weeks."""
    cutoff = time.time() - _CACHE_MAX_AGE_SECONDS
    for path in cache_dir.iterdir():
        try:
            if path.stat().st_mtime < cutoff:
                path.unlink()
        except OSError:
            pass


def install_jinja_cache() -> None:
    """Monkey-patch get_template to use the jinja bytecode cache."""
    _CACHE_DIR.mkdir(parents=True, exist_ok=True)
    _clean_stale_files(_CACHE_DIR)

    import dbt.clients.jinja as dbt_jinja

    try:
        import dbt_common.clients.jinja as dbt_common_jinja
    except ModuleNotFoundError:
        dbt_common_jinja = dbt_jinja  # type: ignore[assignment]  # dbt < 1.8

    _orig_get_template = dbt_common_jinja.get_template
    _code_cache: dict[str, object] = {}

    def _cached_get_template(
        string: str,
        ctx: dict,
        node: object = None,
        capture_macros: bool = False,
        native: bool = False,
    ) -> jinja2.Template:
        if not isinstance(string, str):
            return _orig_get_template(string, ctx, node, capture_macros, native)

        key = _jinja_hash(string)

        if not (code := _code_cache.get(key)):
            path = _CACHE_DIR / f"{key}.marshal"
            try:
                with open(path, "rb") as f:
                    code = marshal.load(f)
                path.touch()
            except Exception:
                pass

        env = dbt_common_jinja.get_environment(node, capture_macros, native=native)

        if not code:
            code = env.compile(str(string))
            try:
                with open(path, "wb") as f:
                    marshal.dump(code, f)
            except Exception:
                pass

        _code_cache[key] = code
        return env.template_class.from_code(env, code, env.make_globals(ctx), None)

    dbt_common_jinja.get_template = _cached_get_template  # type: ignore[assignment]
    dbt_jinja.get_template = _cached_get_template  # type: ignore[assignment]

    # Cache Jinja ASTs so that macro parsing (pass 1) and
    # statically_extract_macro_calls (pass 2) don't re-lex the same source.
    # Safety: the AST is a pure function of (source, lexer, extensions, parser
    # class) — all identical across dbt's Environment instances.  Consumers
    # (find_all, compiler.generate) use a read-only visitor pattern.
    _orig_env_parse = jinja2.Environment.parse
    _ast_cache: dict[str, object] = {}

    def _cached_env_parse(
        self: jinja2.Environment,
        source: str,
        name: str | None = None,
        filename: str | None = None,
    ) -> jinja2.nodes.Template:
        if name is None and filename is None and isinstance(source, str):
            if (ast := _ast_cache.get(source)) is None:
                ast = _ast_cache[source] = _orig_env_parse(self, source, name, filename)
            return ast
        return _orig_env_parse(self, source, name, filename)

    jinja2.Environment.parse = _cached_env_parse  # type: ignore[assignment]
