# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0

"""Dotpromptz is a library for generating prompts using Handlebars templates.

This module provides the core Dotprompt class, which extends the Handlebars
template engine for use with generative AI prompts. It supports parsing
templates, rendering metadata, resolving tools and partials, and managing
schemas.

Key features include:

| Feature              | Description                                                                             |
|----------------------|-----------------------------------------------------------------------------------------|
| Template Parsing     | Parsing of templates with YAML frontmatter for metadata.                                |
| Metadata Rendering   | Rendering of prompt metadata, including merging and resolving tools and schemas.        |
| Tool Resolution      | Resolving tool names to tool definitions using a resolver or a static mapping.          |
| Partial Resolution   | Resolving partial template names to their content using a resolver or a static mapping. |
| Schema Management    | Handling of JSON schemas, including Picoschema conversion.                              |
| Helper Functions     | Registration and management of custom helper functions.                                 |
| Partial Templates    | Registration and management of partial templates.                                       |
| Model Configuration  | Support for default models and model-specific configurations.                           |
| Prompt Store         | Integration with a prompt store for loading prompts and partials.                       |
| Extensibility        | Designed to be extensible with custom helpers, resolvers, and stores.                   |
"""

from __future__ import annotations

import re
from pathlib import Path
from typing import Any

from handlebarrz import Context, EscapeFunction, Handlebars, HelperFn, RuntimeOptions
from handlebarrz import Context, EscapeFunction, Handlebars, HelperFn, RuntimeOptions

from dotpromptz.helpers import BUILTIN_HELPERS
from dotpromptz.input_resolver import resolve_input
from dotpromptz.parse import parse_document, to_messages
from dotpromptz.picoschema import picoschema_to_json_schema
from dotpromptz.resolvers import resolve_json_schema, resolve_partial, resolve_tool
from dotpromptz.typing import (
    DataArgument,
    JsonSchema,
    ModelConfigT,
    ParsedPrompt,
    PartialResolver,
    PromptFunction,
    PromptMetadata,
    PromptMetadata,
    RenderedPrompt,
    SchemaResolver,
    ToolDefinition,
    ToolResolver,
    VariablesT,
)
from dotpromptz.util import remove_undefined_fields

# Pre-compiled regex for finding partial references in handlebars templates

# Since the handlebars-rust implementation doesn't expose a visitor
# to walk the AST to find partial nodes, we're using a crude regular expression
# to find partials.
_PARTIAL_PATTERN = re.compile(r'{{\s*>\s*([a-zA-Z0-9_.-]+)\s*}}')


def _merge_metadata(
    current: PromptMetadata[ModelConfigT],
    merge: PromptMetadata[ModelConfigT],
) -> PromptMetadata[ModelConfigT]:
    """Merges a single metadata object into the current one.

    Args:
        current: The current metadata object.
        merge: The metadata object to merge into the current one.

    Returns:
        The merged metadata object.
    """
    # Convert Pydantic models to raw dicts by alias first. Skip None values.
    merge_dict = merge.model_dump(exclude_none=True, by_alias=True)
    current_dict = current.model_dump(exclude_none=True, by_alias=True)

    # Keep a reference to the original config.
    original_config = current_dict.get('config', {})
    new_config = merge_dict.get('config', {})

    # Merge the new metadata.
    current_dict.update(merge_dict)

    # Merge the configs and set the resulting config.
    current_dict['config'] = {**original_config, **new_config}

    # Recreate the Pydantic model from the merged dict and validate it.
    return PromptMetadata[ModelConfigT].model_validate(current_dict)


def _identify_partials(template: str) -> set[str]:
    """Identify all unique partial references in a template.

    Args:
        template: The template to scan for partial references.

    Returns:
        A set of partial names referenced in the template.
    """
    return set(_PARTIAL_PATTERN.findall(template))


class RenderFunc(PromptFunction[ModelConfigT]):
    """A compiled prompt function with the prompt as a property.

    This is the Python equivalent of the renderFunc nested function
    within the compile method of the Dotprompt class in TypeScript.

    It exposes the prompt property to the user.
    """

    def __init__(
        self,
        dotprompt: Dotprompt,
        handlebars: Handlebars,
        prompt: ParsedPrompt[ModelConfigT],
        prompt_file_path: Path | None = None,
    ):
        """Initialize the renderer.

        Args:
            dotprompt: The Dotprompt instance.
            handlebars: The Handlebars instance.
            prompt: The parsed prompt.
            prompt_file_path: Optional path to the prompt file (for resolving relative input paths).
        """
        self._dotprompt = dotprompt
        self._handlebars = handlebars
        self._prompt_file_path = prompt_file_path if prompt_file_path is not None else Path.cwd()

        self.prompt = prompt

    def __call__(
        self, data: DataArgument[VariablesT] | None = None, options: PromptMetadata[ModelConfigT] | None = None
    ) -> RenderedPrompt[ModelConfigT]:
        """Render the prompt.

        Args:
            data: Optional data to be used to render the prompt. If None, auto-constructs from frontmatter input.
            options: Additional options for the prompt.

        Returns:
            The rendered prompt.
        """
        merged_metadata: PromptMetadata[ModelConfigT] = self._dotprompt.render_metadata(self.prompt, options)

        # Resolve input from frontmatter if needed
        frontmatter_input: dict[str, Any] = {}
        if self.prompt.input is not None and self.prompt.input.data is not None:
            # Resolve frontmatter input via input_resolver
            try:
                records, is_batch = resolve_input(self.prompt.input.data, self._prompt_file_path)
            except ValueError as exc:
                raise ValueError(f'Input resolution failed: {exc}') from exc

            # Batch mode requires caller to handle iteration
            if is_batch:
                raise ValueError(
                    'Batch input detected in frontmatter. Use CLI for batch processing, '
                    'or call resolve_input() manually and iterate over records.'
                )

            # Single mode: use first record (or empty dict if no records)
            frontmatter_input = records[0] if records else {}

        # Build context: merge frontmatter input with caller-provided data (caller overrides)
        if data is None:
            # No caller data → use only frontmatter input
            context: Context = frontmatter_input
        else:
            # Caller provided data → merge (caller overrides frontmatter)
            context = {
                **frontmatter_input,
                **(data.input or {}),
            }

        # Prepare runtime options
        runtime_options: RuntimeOptions = {
            'data': {
                **(data.context if data and data.context else {}),
            },
        }

        # Render the string
        render_string = self._handlebars.compile(self.prompt.template)
        rendered_string = render_string(context, runtime_options)

        # Parse the rendered string into messages
        # Use data if provided, otherwise construct minimal DataArgument with resolved input
        final_data = data if data is not None else DataArgument[VariablesT](input=frontmatter_input)
        messages = to_messages(rendered_string, final_data)

        # Construct and return the final RenderedPrompt
        return RenderedPrompt[ModelConfigT](
            # Spread the metadata fields into the RenderedPrompt constructor.
            **merged_metadata.model_dump(exclude_none=True, by_alias=True),
            messages=messages,
        )


class Dotprompt:
    """Dotprompt extends a Handlebars template for use with Gen AI prompts."""

    def __init__(
        self,
        default_model: str | None = None,
        model_configs: dict[str, Any] | None = None,
        helpers: dict[str, HelperFn] | None = None,
        partials: dict[str, str] | None = None,
        tools: dict[str, ToolDefinition] | None = None,
        tool_resolver: ToolResolver | None = None,
        schemas: dict[str, JsonSchema] | None = None,
        schema_resolver: SchemaResolver | None = None,
        partial_resolver: PartialResolver | None = None,
        escape_fn: EscapeFunction = EscapeFunction.NO_ESCAPE,
    ) -> None:
        """Initialize Dotprompt with a Handlebars template.

        Args:
            default_model: The default model to use for the prompt when not specified in the template.
            model_configs: Assign a set of default configuration options to be used with a particular model.
            helpers: Helpers to pre-register.
            partials: Partials to pre-register.
            tools: Provide a static mapping of tool definitions that should be used when resolving tool names.
            tool_resolver: Provide a lookup implementation to resolve tool names to definitions.
            schemas: Provide a static mapping of schema names to their JSON schema definitions.
            schema_resolver: resolver for schema names to JSON schema definitions.
            partial_resolver: resolver for partial names to their content.
            escape_fn: escape function to use for the template.
        """
        self._handlebars: Handlebars = Handlebars(escape_fn=escape_fn)

        self._known_helpers: dict[str, bool] = {}
        self._default_model: str | None = default_model
        self._model_configs: dict[str, Any] = model_configs or {}
        self._helpers: dict[str, HelperFn] = helpers or {}
        self._partials: dict[str, str] = partials or {}
        self._tools: dict[str, ToolDefinition] = tools or {}
        self._tool_resolver: ToolResolver | None = tool_resolver
        self._schemas: dict[str, JsonSchema] = schemas or {}
        self._schema_resolver: SchemaResolver | None = schema_resolver
        self._partial_resolver: PartialResolver | None = partial_resolver

        self._register_initial_helpers(
            builtin_helpers=BUILTIN_HELPERS,
            custom_helpers=helpers,
        )
        self._register_initial_partials(partials)

    def define_helper(self, name: str, fn: HelperFn) -> Dotprompt:
        """Define a helper function for the template.

        Args:
            name: The name of the helper function.
            fn: The function to be called when the helper is used in the
                template.

        Returns:
            The Dotprompt instance.
        """
        self._handlebars.register_helper(name, fn)
        self._known_helpers[name] = True
        return self

    def define_partial(self, name: str, source: str) -> Dotprompt:
        """Define a partial template for the template.

        Args:
            name: The name of the partial template.
            source: The source code for the partial.

        Returns:
            The Dotprompt instance.
        """
        self._handlebars.register_partial(name, source)
        return self

    def define_tool(self, definition: ToolDefinition) -> Dotprompt:
        """Define a tool for the template.

        Args:
            definition: The definition of the tool.

        Returns:
            The Dotprompt instance.
        """
        self._tools[definition.name] = definition
        return self

    def parse(self, source: str) -> ParsedPrompt[ModelConfigT]:
        """Parse a prompt from a string.

        Args:
            source: The source code for the prompt.

        Returns:
            The parsed prompt.
        """
        return parse_document(source)

    def render(
        self,
        source: str,
        data: DataArgument[VariablesT] | None = None,
        options: PromptMetadata[ModelConfigT] | None = None,
        prompt_file_path: Path | None = None
    ) -> RenderedPrompt[ModelConfigT]:
        """Render a prompt.

        Args:
            source: The source code for the prompt.
            data: Optional data to be used to render the prompt. If None, auto-constructs from frontmatter input.
            options: Additional options for the prompt.
            prompt_file_path: Optional path to the prompt file (for resolving relative input file paths).
                Defaults to cwd.

        Returns:
            The rendered prompt.
        """
        renderer: PromptFunction[ModelConfigT] = self.compile(source, prompt_file_path=prompt_file_path)
        return renderer(data, options)

    def compile(
        self,
        source: str | ParsedPrompt[ModelConfigT],
        additional_metadata: PromptMetadata[ModelConfigT] | None = None,
        prompt_file_path: Path | None = None
    ) -> PromptFunction[ModelConfigT]:
        """Compile a prompt.

        Args:
            source: The source code for the prompt, or an already-parsed
                ``ParsedPrompt`` object.
            additional_metadata: Additional metadata to be used to render the prompt.
            prompt_file_path: Optional path to the prompt file (for resolving relative input file paths).

        Returns:
            A function that can be used to render the prompt.
        """
        prompt: ParsedPrompt[ModelConfigT] = self.parse(source) if isinstance(source, str) else source
        if additional_metadata is not None:
            prompt = prompt.model_copy(
                deep=True,
                update=additional_metadata.model_dump(exclude_none=True, by_alias=True),
            )

        # Resolve partials before compiling.
        self._resolve_partials(prompt.template)
        return RenderFunc(self, self._handlebars, prompt, prompt_file_path)

    def render_metadata(
        self,
        source: str | ParsedPrompt[ModelConfigT],
        additional_metadata: PromptMetadata[ModelConfigT] | None = None,
    ) -> PromptMetadata[ModelConfigT]:
        """Render metadata for a prompt.

        Args:
            source: The source code for the prompt or a parsed prompt.
            additional_metadata: Additional metadata to be used to render the prompt.

        Returns:
            The rendered metadata.
        """
        prompt = self.parse(source) if isinstance(source, str) else source

        # Resolve the effective model name.
        # Priority: additional_metadata config.model > prompt config.model > default_model
        model = self._resolve_model_name(prompt, additional_metadata)

        config: ModelConfigT | None = None
        if model is not None and self._model_configs.get(model) is not None:
            config = self._model_configs.get(model)

        return self._resolve_metadata(
            PromptMetadata[ModelConfigT](
                config=config,
            )
            if config is not None
            else PromptMetadata[ModelConfigT](),
            prompt,
            additional_metadata,
        )

    @staticmethod
    def _resolve_model_name(
        prompt: ParsedPrompt[ModelConfigT],
        additional_metadata: PromptMetadata[ModelConfigT] | None = None,
    ) -> str | None:
        """Resolve the effective LLM model name.

        Priority (highest to lowest):
        1. ``additional_metadata.config['model']`` (runtime override)
        2. ``prompt.config['model']``               (frontmatter config.model)
        3. ``None``                                  (will fall back to Dotprompt._default_model externally)

        Args:
            prompt: The parsed prompt.
            additional_metadata: Optional runtime metadata override.

        Returns:
            The resolved model name, or ``None``.
        """

        def _config_model(meta: PromptMetadata[ModelConfigT] | None) -> str | None:
            if meta is None or meta.config is None:
                return None
            cfg = meta.config
            if isinstance(cfg, dict):
                return cfg.get('model')
            return None

        # Prefer config.model from additional_metadata, then from prompt.
        return _config_model(additional_metadata) or _config_model(prompt)

    def _resolve_metadata(
        self, base: PromptMetadata[ModelConfigT], *merges: PromptMetadata[ModelConfigT] | None
    ) -> PromptMetadata[ModelConfigT]:
        """Merges multiple metadata objects, resolving tools and schemas.

        Later metadata objects override earlier ones.

        Args:
            base: The base metadata object.
            merges: Additional metadata objects to merge into base.

        Returns:
            Merged metadata.
        """
        out = base.model_copy(deep=True)

        for merge in merges:
            if merge:
                out = _merge_metadata(out, merge)

        # If the merged result is a ParsedPrompt (has a 'template' field), strip it
        # so that RenderedPrompt / PromptMetadata consumers don't see it.
        # Use model_copy with exclude instead of delattr, which is unsafe on Pydantic models.
        if isinstance(out, ParsedPrompt):
            out = PromptMetadata[ModelConfigT].model_validate(
                out.model_dump(exclude={'template'}, exclude_none=True, by_alias=True)
            )

        out = remove_undefined_fields(out)
        # Sequential: _render_picoschema depends on _resolve_tools output.
        out = self._resolve_tools(out)
        out = self._render_picoschema(out)
        return out

    def _render_picoschema(self, meta: PromptMetadata[ModelConfigT]) -> PromptMetadata[ModelConfigT]:
        """Render a Picoschema prompt.

        Args:
            meta: The prompt metadata.

        Returns:
            The rendered prompt metadata.
        """
        needs_output_processing = meta.output is not None and meta.output.schema is not None

        if not needs_output_processing:
            return meta

        new_meta = meta.model_copy(deep=True)

        if needs_output_processing and meta.output is not None:
            new_meta.output.schema = picoschema_to_json_schema(
                meta.output.schema,
                self._wrapped_schema_resolver,
            )
        return new_meta

    def _wrapped_schema_resolver(self, name: str) -> JsonSchema | None:
        """Resolve a schema from either instance local mapping or the resolver.

        Args:
            name: The name of the schema to resolve.

        Returns:
            The resolved schema or None if it is not found.
        """
        if name in self._schemas:
            return self._schemas[name]

        if self._schema_resolver is None:
            return None

        # Cache resolved schemas to avoid redundant resolver calls.
        resolved = resolve_json_schema(name, self._schema_resolver)
        if resolved is not None:
            self._schemas[name] = resolved
        return resolved

    def _resolve_tools(self, metadata: PromptMetadata[ModelConfigT]) -> PromptMetadata[ModelConfigT]:
        """Resolve all tools in a prompt.

        Args:
            metadata: The prompt metadata.

        Returns:
            A copy of the prompt metadata with the tools resolved.

        Raises:
            ToolNotFoundError: If a tool is not found in the resolver.
            ToolResolverFailedError: If a tool resolver fails.
            TypeError: If a tool resolver returns an invalid type.
            ValueError: If a tool resolver is not defined.
        """
        out: PromptMetadata[ModelConfigT] = metadata.model_copy(deep=True)
        if out.tools is None:
            return out

        # Resolve tools that are already registered into toolDefs, leave
        # unregistered tools alone.
        unregistered_names: list[str] = []
        out.tool_defs = out.tool_defs or []

        # Collect all the tools:
        # 1. Already registered tools go into toolDefs.
        # 2. If we have a tool resolver, add the names to the list to resolve.
        # 3. Otherwise, add the names to the list of unregistered tools.
        to_resolve: list[str] = []
        have_resolver = self._tool_resolver is not None
        for name in out.tools:
            if name in self._tools:
                # Found locally.
                out.tool_defs.append(self._tools[name])
            elif have_resolver:
                # Resolve using the tool resolver.
                to_resolve.append(name)
            else:
                # Unregistered tool.
                unregistered_names.append(name)

        # Resolve all the tools using the resolver.
        for tool_name in to_resolve:
            tool = resolve_tool(tool_name, self._tool_resolver)
            if out.tool_defs is not None:
                out.tool_defs.append(tool)

        out.tools = unregistered_names
        return out

    def _resolve_partials(
        self,
        template: str,
        visited: set[str] | None = None,
    ) -> None:
        """Resolve all partials in a template.

        This method recursively resolves partials, meaning if a partial itself
        contains partial references, those will also be resolved. Cycle detection
        prevents infinite loops when partials reference each other.

        Args:
            template: The template to resolve partials in.
            visited: Set of partial names currently being processed (for cycle detection).

        Returns:
            None
        """
        if self._partial_resolver is None:
            return

        if visited is None:
            visited = set()

        names = _identify_partials(template)

        for name in names:
            if not self._handlebars.has_partial(name) and name not in visited:
                visited.add(name)
                content: str | None = None

                if self._partial_resolver is not None:
                    content = resolve_partial(name, self._partial_resolver)

                if content is not None:
                    self.define_partial(name, content)
                    # Recursively resolve partials in the content.
                    self._resolve_partials(content, visited)
    def _register_initial_helpers(
        self,
        builtin_helpers: dict[str, HelperFn] | None = None,
        custom_helpers: dict[str, HelperFn] | None = None,
    ) -> None:
        """Register the initial helpers.

        Built-in helpers are registered first, then custom helpers are
        registered.

        Args:
            builtin_helpers: Built-in helpers to register.
            custom_helpers: Custom helpers to register.
        """
        if builtin_helpers is not None:
            for name, fn in builtin_helpers.items():
                self.define_helper(name, fn)

        if custom_helpers is not None:
            for name, fn in custom_helpers.items():
                self.define_helper(name, fn)

    def _register_initial_partials(self, partials: dict[str, str] | None = None) -> None:
        """Register the initial partials.

        Args:
            partials: Partials to register.
        """
        if partials is not None:
            for name, source in partials.items():
                self.define_partial(name, source)
