"""Jinja2-based prompt template engine."""

from __future__ import annotations

from pathlib import Path
from typing import Any

try:
    from jinja2 import Environment, FileSystemLoader, StrictUndefined, TemplateNotFound, UndefinedError
except ImportError as e:
    raise ImportError("Jinja2 not installed. Install with: pip install koan[prompts]") from e

_BUILTIN_TEMPLATES = Path(__file__).parent / "templates"


class PromptRenderError(Exception):
    """Raised when a prompt template cannot be rendered."""

    def __init__(self, message: str, template_name: str | None = None, cause: Exception | None = None) -> None:
        self.template_name = template_name
        self.cause = cause
        super().__init__(message)


class PromptEngine:
    """Jinja2 prompt template engine.

    Loads templates from one or more directories. Built-in KOAN templates
    are always available as a fallback.

    Usage::

        engine = PromptEngine()
        prompt = engine.render("system.j2", role="analyst", language="English")

        # Or render from a raw string:
        prompt = engine.render_string("Hello {{ name }}", name="World")
    """

    def __init__(self, template_dirs: list[str | Path] | None = None) -> None:
        search_paths = [str(p) for p in (template_dirs or [])]
        search_paths.append(str(_BUILTIN_TEMPLATES))

        self._env = Environment(
            loader=FileSystemLoader(search_paths),
            undefined=StrictUndefined,
            trim_blocks=True,
            lstrip_blocks=True,
            keep_trailing_newline=False,
        )

    def render(self, template_name: str, **variables: Any) -> str:
        """Render a named template file with the given variables.

        Args:
            template_name: Template filename (e.g., "system.j2").
            **variables: Template variables.

        Returns:
            Rendered prompt string.

        Raises:
            PromptRenderError: If template not found or variable missing.
        """
        try:
            template = self._env.get_template(template_name)
            return template.render(**variables)
        except TemplateNotFound as e:
            raise PromptRenderError(
                f"Template not found: {template_name}",
                template_name=template_name,
                cause=e,
            ) from e
        except UndefinedError as e:
            raise PromptRenderError(
                f"Missing variable in template '{template_name}': {e}",
                template_name=template_name,
                cause=e,
            ) from e

    def render_string(self, template_str: str, **variables: Any) -> str:
        """Render a template from a raw string.

        Args:
            template_str: Jinja2 template string.
            **variables: Template variables.

        Returns:
            Rendered prompt string.

        Raises:
            PromptRenderError: If a required variable is missing.
        """
        try:
            template = self._env.from_string(template_str)
            return template.render(**variables)
        except UndefinedError as e:
            raise PromptRenderError(
                f"Missing variable in inline template: {e}",
                cause=e,
            ) from e
