"""
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

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 _detect_circular_references(plan: ExecutionPlan) -> None:
    """Raise ``ValueError`` if any steps form a circular dependency.

    Builds a directed graph from each step's ``context_refs`` and performs a
    DFS-based cycle check.  Only refs that point to other steps (i.e. entries
    matching an existing ``step_id``) are considered edges.
    """
    step_ids = {s.step_id for s in plan.steps}

    # adjacency: step_id -> set of step_ids it depends on
    graph: Dict[str, set] = {s.step_id: set() for s in plan.steps}
    for s in plan.steps:
        for ref in s.context_refs:
            if ref in step_ids and ref != s.step_id:
                graph[s.step_id].add(ref)

    # DFS cycle detection
    WHITE, GRAY, BLACK = 0, 1, 2
    color: Dict[str, int] = {sid: WHITE for sid in step_ids}

    def _dfs(node: str, path: list) -> None:
        color[node] = GRAY
        path.append(node)
        for neighbour in graph.get(node, set()):
            if color[neighbour] == GRAY:
                # Found a cycle — extract the loop portion of the path
                cycle_start = path.index(neighbour)
                cycle = path[cycle_start:] + [neighbour]
                raise ValueError(
                    f"Circular reference detected: {' -> '.join(cycle)}"
                )
            if color[neighbour] == WHITE:
                _dfs(neighbour, path)
        path.pop()
        color[node] = BLACK

    for sid in step_ids:
        if color[sid] == WHITE:
            _dfs(sid, [])


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.
    """
    _detect_circular_references(plan)

    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)

    # Replace any remaining unresolved step references with a clear marker
    remaining = _REF_PATTERN.findall(resolved)
    for ref_name in remaining:
        if ref_name.startswith("step_") and ref_name.endswith("_output"):
            # Extract step id (e.g. "step_1" from "step_1_output")
            referenced_step = ref_name.removesuffix("_output")
            if referenced_step not in previous_results:
                placeholder = "{" + ref_name + "}"
                available = sorted(previous_results.keys()) if previous_results else []
                available_refs = (
                    ", ".join(f"{{{s}_output}}" for s in available)
                    if available
                    else "(none)"
                )
                replacement = f"[{referenced_step} not found]"
                resolved = resolved.replace(placeholder, replacement)
                log.warning(
                    "Context ref '{%s}' in step '%s' references non-existent step '%s' — "
                    "substituting '%s'. Available refs: %s",
                    ref_name,
                    step.step_id,
                    referenced_step,
                    replacement,
                    available_refs,
                )

    # Context size estimation
    estimated_tokens = len(resolved) // 4
    if estimated_tokens > 50000:
        log.warning(
            "Large context for step '%s': ~%d estimated tokens (%d chars)",
            step.step_id,
            estimated_tokens,
            len(resolved),
        )

    return resolved
