"""Auto-variable registry and render-time injection helpers."""

from __future__ import annotations

import json
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any

import jinja2

from dotpromptz.typing import ParsedPrompt

AutoVarProvider = Callable[['AutoVarContext'], Any | None]


@dataclass(frozen=True)
class AutoVarContext:
    """Context data used by auto-variable providers."""

    prompt: ParsedPrompt | None = None
    prompt_file_path: Path | None = None
    now: datetime | None = None
    chain_source: Path | None = None
    chain_text: str | None = None


def _time_auto_var(ctx: AutoVarContext) -> str:
    """Provide ``TIME`` in ``YYYYMMDD_HHMMSS`` format."""
    ts = ctx.now or datetime.now()
    return ts.strftime('%Y%m%d_%H%M%S')


def _name_auto_var(ctx: AutoVarContext) -> str | None:
    """Provide ``NAME`` from the prompt file stem."""
    prompt_path = ctx.prompt_file_path
    if prompt_path is None:
        return None
    p = Path(prompt_path) if not isinstance(prompt_path, Path) else prompt_path
    if not p.is_file():
        return None
    return p.stem


def _schema_auto_var(ctx: AutoVarContext) -> str | None:
    """Provide ``SCHEMA`` from ``output.schema`` as JSON text."""
    if ctx.prompt is None:
        return None
    output_cfg = ctx.prompt.output
    if output_cfg is None:
        return None
    if isinstance(output_cfg, dict):
        schema_val = output_cfg.get('schema')
    else:
        schema_val = getattr(output_cfg, 'schema_value', None)
    if schema_val is None or not isinstance(schema_val, dict):
        return None
    return json.dumps(schema_val, indent=2, ensure_ascii=False)


def _index_auto_var(_: AutoVarContext) -> None:
    """Reserve ``INDEX`` name in the auto-variable registry."""
    return None


def _mumber_auto_var(_: AutoVarContext) -> str:
    """Provide default ``NUMBER`` value for output collision resolution."""
    return '1'


def _text_auto_var(ctx: AutoVarContext) -> str | None:
    """Provide ``TEXT`` from a chain-input text file."""
    if ctx.chain_source is None:
        return None
    if ctx.chain_source.suffix.lower() != '.txt':
        return None
    return ctx.chain_text


AUTO_VAR_PROVIDERS: dict[str, AutoVarProvider] = {
    'TIME': _time_auto_var,
    'NAME': _name_auto_var,
    'SCHEMA': _schema_auto_var,
    'INDEX': _index_auto_var,
    'NUMBER': _mumber_auto_var,
    'TEXT': _text_auto_var,
}
"""Registered auto-variable providers keyed by variable name."""


def auto_var_names() -> frozenset[str]:
    """Return registered auto-variable names."""
    return frozenset(AUTO_VAR_PROVIDERS.keys())


def compute_auto_vars(ctx: AutoVarContext) -> dict[str, Any]:
    """Compute auto-variable values for one render invocation.

    Args:
        ctx: Auto-variable context.

    Returns:
        Mapping of auto-variable names to resolved values.
    """
    resolved: dict[str, Any] = {}
    for name, provider in AUTO_VAR_PROVIDERS.items():
        value = provider(ctx)
        if value is not None:
            resolved[name] = value
    return resolved


def build_render_env(
    base_env: jinja2.Environment,
    auto_globals: Mapping[str, Any] | None = None,
) -> jinja2.Environment:
    """Create a render-scoped Jinja2 overlay environment.

    Args:
        base_env: Shared base Jinja environment.
        auto_globals: Auto variables to inject as globals.

    Returns:
        A per-render Jinja2 environment.
    """
    render_env = base_env.overlay()
    # ``overlay()`` shares the globals mapping with the parent environment.
    # Copy it so per-render globals do not leak across renders.
    render_env.globals = dict(base_env.globals)
    if auto_globals:
        render_env.globals.update(dict(auto_globals))
    return render_env
