"""
Context Resolver — Builds final prompts by substituting context references.

Takes a step's prompt template and replaces placeholders like {original_input}
and {step_1_output} with actual content from previous step results.

Unknown placeholders are left as-is rather than raising errors, so downstream
prompts remain readable even if an upstream step failed.
"""

from __future__ import annotations

import logging
import re
from typing import TYPE_CHECKING, Dict, Optional

if TYPE_CHECKING:
    from .plan_executor import ExecutionPlan, ExecutionStep, StepResult

log = logging.getLogger("brynq.runtime.context")

# Regex to find all {placeholder} refs in a template.
_REF_PATTERN = re.compile(r"\{([a-zA-Z_][a-zA-Z0-9_]*)\}")


def resolve_context(
    step: ExecutionStep,
    plan: ExecutionPlan,
    previous_results: Dict[str, StepResult],
) -> str:
    """Build the full prompt for a step by resolving context references.

    Substitution rules:
        - ``{original_input}`` is replaced with ``plan.original_input``.
        - ``{step_N_output}`` is replaced with the output of step N
          (1-indexed, matching the step's ``step_id`` field format).
        - Any reference to a step that failed or has no output resolves to
          an empty string, with a warning logged.
        - Unrecognised placeholders are left untouched.

    Args:
        step: The current step whose prompt template needs resolving.
        plan: The full execution plan (provides ``original_input``).
        previous_results: Map of ``step_id`` -> ``StepResult`` for all
            previously executed steps.

    Returns:
        The fully resolved prompt string, ready to send to a bridge.
    """
    template = step.prompt

    # Build the substitution map
    subs: Dict[str, str] = {
        "original_input": plan.original_input,
    }

    # Add all previous step outputs
    for step_id, result in previous_results.items():
        key = f"{step_id}_output"
        if result.success:
            subs[key] = result.output
        else:
            subs[key] = ""
            log.warning(
                "Context ref '%s' points to failed step '%s' — substituting empty string",
                key,
                step_id,
            )

    # Also resolve explicit context_refs listed on the step
    for ref in step.context_refs:
        if ref == "original_input":
            continue  # Already in subs
        # Refs like "step_1" should map to key "step_1_output"
        output_key = f"{ref}_output"
        if output_key not in subs and ref in previous_results:
            result = previous_results[ref]
            subs[output_key] = result.output if result.success else ""

    # Perform substitutions
    resolved = template
    for key, value in subs.items():
        placeholder = "{" + key + "}"
        resolved = resolved.replace(placeholder, value)

    # Warn about any remaining unresolved refs that look like step references
    remaining = _REF_PATTERN.findall(resolved)
    for ref_name in remaining:
        if ref_name.startswith("step_") and ref_name.endswith("_output"):
            log.warning(
                "Unresolved context ref '{%s}' in step '%s' — "
                "referenced step may not have executed yet",
                ref_name,
                step.step_id,
            )

    return resolved
