"""Compile and render services for ``.prompt`` files."""

from __future__ import annotations

import re
from collections.abc import Callable, Iterable, Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import jinja2
from jinja2 import meta as jinja2_meta

from dotpromptz.autovars import AutoVarContext, auto_var_names, build_render_env, compute_auto_vars
from dotpromptz.blocks import parse_blocks_to_messages
from dotpromptz.errors import (
    FrontmatterValidationError,
    InputResolutionError,
    InputVariableNamingError,
    PromptParseError,
    TemplateMissingVariableError,
    TemplateVariableNamingError,
)
from dotpromptz.image import normalize_media_messages
from dotpromptz.inputs import resolve_files, resolve_input
from dotpromptz.logging import get_logger
from dotpromptz.parser import parse_document
from dotpromptz.typing import AdapterConfig, ParsedPrompt, PromptOutputConfig, RenderedPrompt, RuntimeConfig

logger = get_logger(__name__)

_AUTO_VARIABLE_PATTERN = re.compile(r'^[A-Z][A-Z0-9_]*$')
_HELPER_NAME_PATTERN = re.compile(r'^[a-z][a-z0-9_]*$')


@dataclass(frozen=True)
class CompiledPrompt:
    """Compiled prompt with resolved input records and render metadata."""

    prompt_path: Path
    source: str
    parsed_prompt: ParsedPrompt
    input_list: list[dict[str, Any]]
    body_required_vars: frozenset[str]
    frontmatter_vars: frozenset[str]
    env: jinja2.Environment
    default_model: str | None
    model_configs: dict[str, Any]
    helpers: dict[str, Callable[..., str]] | None
    chain_source: Path | None = None
    chain_text: str | None = None


@dataclass(frozen=True)
class ResolvedPromptInput:
    """Resolved input records and optional chain-source metadata."""

    input_list: list[dict[str, Any]]
    chain_source: Path | None = None
    chain_text: str | None = None


def _create_jinja_env() -> jinja2.Environment:
    """Create a pre-configured Jinja2 environment for prompt rendering."""
    return jinja2.Environment(
        autoescape=False,
        undefined=jinja2.ChainableUndefined,
        keep_trailing_newline=True,
        trim_blocks=True,
        lstrip_blocks=True,
    )


def _is_auto_variable_name(name: str) -> bool:
    """Return whether ``name`` belongs to the ALL_CAPS auto-variable namespace."""
    return _AUTO_VARIABLE_PATTERN.fullmatch(name) is not None


def _validate_input_variable_keys(keys: Iterable[str]) -> None:
    """Validate input keys against ALL_CAPS auto-variable reservation."""
    invalid = {name for name in keys if _is_auto_variable_name(name)}
    if invalid:
        raise InputVariableNamingError(invalid)


def _validate_template_variable_names(
    variable_names: Iterable[str],
    *,
    auto_names: frozenset[str],
) -> None:
    """Validate template variables against auto-variable registry."""
    unknown_auto = {name for name in variable_names if _is_auto_variable_name(name) and name not in auto_names}
    if unknown_auto:
        raise TemplateVariableNamingError(
            unknown_auto=unknown_auto,
            allowed_auto=set(auto_names),
        )


def _extract_template_variables(template: str, env: jinja2.Environment) -> set[str]:
    """Extract undeclared variables from a Jinja2 template string."""
    try:
        ast = env.parse(template)
    except jinja2.TemplateSyntaxError:
        return set()
    raw_vars = jinja2_meta.find_undeclared_variables(ast)
    excluded = set(env.globals.keys())
    return {name for name in raw_vars if name not in excluded}


def _extract_frontmatter_template_variables(
    frontmatter: dict[str, Any],
    env: jinja2.Environment,
    *,
    exclude_keys: set[str] | None = None,
) -> set[str]:
    """Extract template variable names referenced in frontmatter strings."""
    if exclude_keys is None:
        exclude_keys = {'input'}

    def _walk(value: Any) -> set[str]:
        if isinstance(value, str):
            return _extract_template_variables(value, env)
        if isinstance(value, dict):
            names: set[str] = set()
            for nested in value.values():
                names.update(_walk(nested))
            return names
        if isinstance(value, list):
            names: set[str] = set()
            for item in value:
                names.update(_walk(item))
            return names
        return set()

    names: set[str] = set()
    for key, value in frontmatter.items():
        if key in exclude_keys:
            continue
        names.update(_walk(value))
    return names


def _render_frontmatter_values(
    data: dict[str, Any],
    variables: dict[str, Any],
    *,
    exclude_keys: set[str] | None = None,
    auto_globals: Mapping[str, Any] | None = None,
) -> dict[str, Any]:
    """Recursively render string leaf values in frontmatter dictionaries."""
    if exclude_keys is None:
        exclude_keys = {'input'}

    env = jinja2.Environment(autoescape=False, undefined=jinja2.Undefined)
    if auto_globals:
        env.globals.update(dict(auto_globals))

    def _render_value(value: Any) -> Any:
        if isinstance(value, str):
            return env.from_string(value).render(variables)
        if isinstance(value, dict):
            return {k: _render_value(v) for k, v in value.items()}
        if isinstance(value, list):
            return [_render_value(item) for item in value]
        return value

    return {key: (_render_value(value) if key not in exclude_keys else value) for key, value in data.items()}


def _normalize_output_for_revalidation(raw_output: dict[str, Any]) -> dict[str, Any]:
    """Normalize output payload before model re-validation.

    During initial parse, ``output.example`` is expanded into ``output.schema``.
    A later render pass re-validates frontmatter values. If both keys are still
    present, that second pass would falsely look like a user provided both.
    """
    normalized = dict(raw_output)
    if normalized.get('example') is not None and normalized.get('schema') is not None:
        normalized.pop('schema', None)
    return normalized


def _analyze_prompt_variables(prompt: ParsedPrompt, env: jinja2.Environment) -> tuple[frozenset[str], frozenset[str]]:
    """Analyze required template variables in body and frontmatter."""
    body_vars = _extract_template_variables(prompt.template or '', env)
    frontmatter_vars = _extract_frontmatter_template_variables(
        prompt.model_dump(by_alias=True, exclude_none=True, exclude={'template', 'warnings'}),
        env,
    )

    auto_names = auto_var_names()
    _validate_template_variable_names(body_vars | frontmatter_vars, auto_names=auto_names)

    body_required_vars = frozenset(name for name in body_vars if name not in auto_names)
    frontmatter_used_vars = frozenset(name for name in frontmatter_vars if name not in auto_names)
    return body_required_vars, frontmatter_used_vars


def _re_render_frontmatter(
    prompt: ParsedPrompt,
    context: dict[str, Any],
    auto_globals: Mapping[str, Any],
) -> ParsedPrompt:
    """Re-render frontmatter string fields against the current input context."""
    rendered_raw = _render_frontmatter_values(
        prompt.model_dump(by_alias=True, exclude_none=False, exclude={'template', 'warnings'}),
        context,
        auto_globals=auto_globals,
    )

    update: dict[str, Any] = {}
    for key in ('description', 'adapter', 'config', 'output', 'runtime'):
        if key in rendered_raw:
            update[key] = rendered_raw[key]

    if not update:
        return prompt

    if isinstance(update.get('adapter'), dict):
        update['adapter'] = AdapterConfig.model_validate(update['adapter'])
    if isinstance(update.get('output'), dict):
        update['output'] = PromptOutputConfig.model_validate(
            _normalize_output_for_revalidation(update['output']),
        )
    if isinstance(update.get('runtime'), dict):
        update['runtime'] = RuntimeConfig.model_validate(update['runtime'])

    payload = prompt.model_dump(by_alias=True, exclude_none=False)
    payload.update(update)
    return ParsedPrompt.model_validate(payload)


def _apply_model_defaults(
    metadata: dict[str, Any],
    *,
    default_model: str | None,
    model_configs: dict[str, Any],
) -> dict[str, Any]:
    """Apply default model and model-specific config defaults."""
    config_raw = metadata.get('config')
    config = dict(config_raw) if isinstance(config_raw, dict) else {}

    model_name = config.get('model')
    if isinstance(model_name, str) and model_name:
        model_defaults = model_configs.get(model_name)
        if isinstance(model_defaults, dict):
            config = {**model_defaults, **config}
    elif default_model:
        model_defaults = model_configs.get(default_model)
        merged_defaults = dict(model_defaults) if isinstance(model_defaults, dict) else {}
        config = {**merged_defaults, **config, 'model': default_model}

    if config:
        metadata['config'] = config
    return metadata


def _parse_prompt_source(source: str) -> ParsedPrompt:
    """Parse raw prompt source into ``ParsedPrompt``."""
    try:
        return parse_document(source)
    except PromptParseError:
        raise
    except ValueError as exc:
        raise FrontmatterValidationError(str(exc), code='frontmatter_invalid') from exc


def _build_render_env_with_helpers(helpers: dict[str, Callable[..., str]] | None = None) -> jinja2.Environment:
    """Create render environment and register optional helper globals."""
    env = _create_jinja_env()
    if helpers:
        for name, fn in helpers.items():
            if _HELPER_NAME_PATTERN.fullmatch(name) is None:
                raise ValueError(f'Helper names must match ^[a-z][a-z0-9_]*$ (got {name!r}).')
            env.globals[name] = fn
    return env


def _build_input_list(
    records: list[dict[str, Any]],
    *,
    defaults: dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
    """Normalize resolved records into the compiled input list shape."""
    if not records and defaults:
        return [defaults]
    if not records:
        return [{}]
    return records


def _resolve_chain_source_path(chain_source: Path | None) -> Path:
    """Resolve and validate a runtime-provided chain source path."""
    if chain_source is None:
        raise ValueError('input.chain requires a runtime-provided chain source.')

    resolved = Path(chain_source).resolve()
    if not resolved.is_file():
        raise ValueError(f'Chain input source not found: {chain_source}')
    return resolved


def _resolve_prompt_input(
    prompt: ParsedPrompt,
    *,
    prompt_path: Path,
    chain_source: Path | None = None,
    single_render: bool = False,
) -> ResolvedPromptInput:
    """Resolve prompt input records for compile/render operations."""
    input_cfg = prompt.input
    if input_cfg is None:
        return ResolvedPromptInput(input_list=[{}])

    defaults = input_cfg.defaults

    if input_cfg.chain:
        resolved_chain_source = _resolve_chain_source_path(chain_source)
        suffix = resolved_chain_source.suffix.lower()
        if suffix in {'.json', '.yaml', '.yml', '.jsonl'}:
            records = resolve_files(str(resolved_chain_source), prompt_path, defaults=defaults)
            input_list = _build_input_list(records, defaults=defaults)
            if single_render and len(input_list) > 1:
                raise ValueError(
                    'Batch input detected in frontmatter. '
                    'Single render does not accept multiple records. '
                    'Use compile_prompt()/plan_prompt() to iterate over records.',
                )
            return ResolvedPromptInput(
                input_list=input_list,
                chain_source=resolved_chain_source,
            )

        if suffix == '.txt':
            try:
                chain_text = resolved_chain_source.read_text(encoding='utf-8')
            except OSError as exc:
                raise ValueError(f'Failed to read chain input source {resolved_chain_source}: {exc}') from exc

            return ResolvedPromptInput(
                input_list=_build_input_list([], defaults=defaults),
                chain_source=resolved_chain_source,
                chain_text=chain_text,
            )

        if suffix == '.png':
            raise ValueError('input.chain does not support image outputs yet.')

        raise ValueError(f'input.chain does not support chain source format {suffix or "<none>"}')

    if input_cfg.data is None and input_cfg.files is None and input_cfg.defaults is None:
        return ResolvedPromptInput(input_list=[{}])

    if input_cfg.files is not None:
        if single_render and not prompt_path.is_file():
            raise ValueError(
                (
                    'render_prompt() with frontmatter input.files requires '
                    'prompt_file_path to point to a real .prompt file.'
                ),
            )
        records = resolve_files(input_cfg.files, prompt_path, defaults=defaults)
    else:
        records = resolve_input(input_cfg.data, prompt_path, defaults=defaults)

    input_list = _build_input_list(records, defaults=defaults)
    if single_render and len(input_list) > 1:
        raise ValueError(
            'Batch input detected in frontmatter. '
            'Single render does not accept multiple records. '
            'Use compile_prompt()/plan_prompt() to iterate over records.',
        )
    return ResolvedPromptInput(input_list=input_list)


def _resolve_frontmatter_input_for_render(
    prompt: ParsedPrompt,
    *,
    prompt_path: Path,
    chain_source: Path | None = None,
) -> ResolvedPromptInput:
    """Resolve frontmatter ``input`` for one render invocation."""
    resolved = _resolve_prompt_input(
        prompt,
        prompt_path=prompt_path,
        chain_source=chain_source,
        single_render=True,
    )
    return ResolvedPromptInput(
        input_list=[resolved.input_list[0]],
        chain_source=resolved.chain_source,
        chain_text=resolved.chain_text,
    )


def compile_prompt(
    prompt_file: str | Path,
    *,
    default_model: str | None = None,
    model_configs: dict[str, Any] | None = None,
    helpers: dict[str, Callable[..., str]] | None = None,
    chain_source: Path | None = None,
) -> CompiledPrompt:
    """Compile a ``.prompt`` file for execution.

    Args:
        prompt_file: Path to a ``.prompt`` file.
        default_model: Default model used when prompt config omits model.
        model_configs: Per-model default config map.
        helpers: Optional Jinja helpers registered as template globals.
        chain_source: Optional runtime-provided chain input file path.

    Returns:
        Compiled prompt object.

    Raises:
        InputResolutionError: If prompt file or input resolution fails.
        FrontmatterValidationError: If parsing fails.
    """
    prompt_path = Path(prompt_file)
    if not prompt_path.is_file():
        raise InputResolutionError(
            f'Path does not exist: {prompt_file}',
            code='prompt_file_missing',
        )

    try:
        source = prompt_path.read_text(encoding='utf-8')
    except OSError as exc:
        raise InputResolutionError(
            f'Failed to read prompt file {prompt_path}: {exc}',
            code='prompt_file_read_failed',
        ) from exc

    env = _build_render_env_with_helpers(helpers)

    parsed = _parse_prompt_source(source)
    body_required_vars, frontmatter_vars = _analyze_prompt_variables(parsed, env)

    try:
        resolved_input = _resolve_prompt_input(
            parsed,
            prompt_path=prompt_path,
            chain_source=chain_source,
        )
    except ValueError as exc:
        raise InputResolutionError(
            f'Input resolution failed: {exc}',
            code='input_resolution_failed',
        ) from exc

    return CompiledPrompt(
        prompt_path=prompt_path,
        source=source,
        parsed_prompt=parsed,
        input_list=resolved_input.input_list,
        body_required_vars=body_required_vars,
        frontmatter_vars=frontmatter_vars,
        env=env,
        default_model=default_model,
        model_configs=model_configs or {},
        helpers=helpers,
        chain_source=resolved_input.chain_source,
        chain_text=resolved_input.chain_text,
    )


def render_compiled_prompt(
    compiled: CompiledPrompt,
    input_vars: dict[str, Any],
    *,
    auto_vars_override: Mapping[str, Any] | None = None,
) -> RenderedPrompt:
    """Render a compiled prompt with one input record.

    Args:
        compiled: Compiled prompt object.
        input_vars: Input variables for one rendered item.
        auto_vars_override: Optional per-render auto-variable overrides.

    Returns:
        Rendered prompt payload.

    Raises:
        TemplateMissingVariableError: If required template variables are missing.
        InputVariableNamingError: If input keys use reserved ALL_CAPS names.
    """
    context = dict(input_vars)
    _validate_input_variable_keys(context.keys())

    provided_vars = set(context.keys())
    missing = set(compiled.body_required_vars) - provided_vars
    if missing:
        raise TemplateMissingVariableError(
            missing=missing,
            template_snippet=(compiled.parsed_prompt.template or '')[:80],
        )

    unused = provided_vars - set(compiled.body_required_vars | compiled.frontmatter_vars)
    if unused:
        logger.warning(
            'Data contains unused variable(s) not referenced in template.',
            variables=', '.join(sorted(unused)),
        )

    for variable in sorted(set(compiled.body_required_vars) & provided_vars):
        value = context[variable]
        if value is None or (isinstance(value, (str, list, dict)) and not value):
            logger.warning('Template variable has an empty value.', variable=variable, value=repr(value))

    auto_ctx = AutoVarContext(
        prompt=compiled.parsed_prompt,
        prompt_file_path=compiled.prompt_path,
        chain_source=compiled.chain_source,
        chain_text=compiled.chain_text,
    )
    auto_vars = compute_auto_vars(auto_ctx)
    if auto_vars_override:
        auto_vars.update(dict(auto_vars_override))

    render_env = build_render_env(compiled.env, auto_vars)

    active_prompt = _re_render_frontmatter(compiled.parsed_prompt, context, auto_globals=auto_vars)
    rendered_string = render_env.from_string(compiled.parsed_prompt.template or '').render(context)
    messages = parse_blocks_to_messages(rendered_string)
    messages = normalize_media_messages(messages, prompt_file_path=compiled.prompt_path)

    metadata = active_prompt.model_dump(exclude_none=True, by_alias=True, exclude={'template', 'warnings'})
    metadata = _apply_model_defaults(
        metadata,
        default_model=compiled.default_model,
        model_configs=compiled.model_configs,
    )
    if isinstance(metadata.get('output'), dict):
        metadata['output'] = _normalize_output_for_revalidation(metadata['output'])

    return RenderedPrompt.model_validate(
        {
            **metadata,
            'messages': messages,
        }
    )


def render_prompt(
    source: str | ParsedPrompt,
    *,
    input: dict[str, Any] | None = None,
    prompt_file_path: Path | None = None,
    default_model: str | None = None,
    model_configs: dict[str, Any] | None = None,
    helpers: dict[str, Callable[..., str]] | None = None,
    chain_source: Path | None = None,
) -> RenderedPrompt:
    """Render one prompt source into a ``RenderedPrompt``.

    Args:
        source: Raw ``.prompt`` text or already parsed prompt.
        input: Optional runtime input variables.
        prompt_file_path: Optional real prompt file path (needed when frontmatter uses ``input.files``).
        default_model: Optional default model when config omits ``model``.
        model_configs: Optional model-default config map.
        helpers: Optional helper functions registered as Jinja globals.
        chain_source: Optional runtime-provided chain input file path.

    Returns:
        Rendered prompt payload.
    """
    env = _build_render_env_with_helpers(helpers)
    prompt = _parse_prompt_source(source) if isinstance(source, str) else source
    prompt_path = prompt_file_path or Path.cwd()
    resolved_input = _resolve_frontmatter_input_for_render(
        prompt,
        prompt_path=prompt_path,
        chain_source=chain_source,
    )
    frontmatter_input = resolved_input.input_list[0]
    context = {**frontmatter_input, **(input or {})}
    body_required_vars, frontmatter_vars = _analyze_prompt_variables(prompt, env)
    compiled = CompiledPrompt(
        prompt_path=prompt_path,
        source=source if isinstance(source, str) else '',
        parsed_prompt=prompt,
        input_list=[context],
        body_required_vars=body_required_vars,
        frontmatter_vars=frontmatter_vars,
        env=env,
        default_model=default_model,
        model_configs=model_configs or {},
        helpers=helpers,
        chain_source=resolved_input.chain_source,
        chain_text=resolved_input.chain_text,
    )
    return render_compiled_prompt(compiled, context)
