"""Build the PreToolUse prompt from session state + proposed action.

The <signals> block surfaces the actual decision features as explicit
booleans so the model latches onto them instead of re-deriving from raw
state. Both the synthetic and real builders call format_prompt, so train
and inference prompts stay identical.
"""
import re

TEST_RE = re.compile(r"\b(pytest|npm test|pnpm test|yarn test|cargo test|go test|bun test|unittest|jest|vitest|tox)\b")

TEMPLATE = """<task>
{goal}
</task>

<recent_trajectory>
{trajectory}
</recent_trajectory>

<state>
observed_files: {observed_files}
edited_files: {edited_files}
traceback_path: {traceback_path}
traceback_symbols: {traceback_symbols}
last_test_status: {last_test_status}
last_search_result_class: {last_search_result_class}
code_changed_since_last_test: {code_changed_since_last_test}
</state>

<proposed_action>
Tool: {tool}
Args: {args}
</proposed_action>

<signals>
proposed_edit_path: {proposed_edit_path}
proposed_action_is_test: {proposed_action_is_test}
traceback_file_read: {traceback_file_read}
edit_target_read: {edit_target_read}
last_failure_type: {last_failure_type}
repeated_command: {repeated_command}
</signals>

<instruction>
Output exactly one short sentence of guidance, or NO_HINT.
</instruction>
"""


def _join(v, last=None):
    if not v:
        return "none"
    if isinstance(v, (list, tuple)):
        v = list(v)[-last:] if last else list(v)
        return ", ".join(str(x) for x in v) if v else "none"
    return str(v)


def _b(x):
    return str(bool(x)).lower()


def format_prompt(goal, trajectory, state, proposed):
    goal = goal.strip()
    if len(goal) > 500:
        goal = goal[:500] + " …"
    tool = proposed["name"]
    args = proposed.get("args", {}) or {}
    cmd = args.get("command", "") or ""
    edit_path = args.get("file_path") if tool in {"Edit", "Write", "MultiEdit", "NotebookEdit"} else None
    observed = set(state.get("observed_files") or [])
    tb = (state.get("traceback_paths") or [None])[-1]
    hist = state.get("command_history") or []

    return TEMPLATE.format(
        goal=goal,
        trajectory=trajectory.strip() if isinstance(trajectory, str) else _fmt_traj(trajectory),
        observed_files=_join(state.get("observed_files"), last=20),
        edited_files=_join(state.get("edited_files"), last=20),
        traceback_path=tb or "none",
        traceback_symbols=_join(state.get("traceback_symbols"), last=5),
        last_test_status=state.get("last_test_status", "unknown"),
        last_search_result_class=state.get("last_search_result_class") or "none",
        code_changed_since_last_test=_b(state.get("code_changed_since_last_test")),
        tool=tool,
        args=_join(edit_path or cmd or args)[:200],
        proposed_edit_path=edit_path or "none",
        proposed_action_is_test=_b(tool == "Bash" and TEST_RE.search(cmd)),
        traceback_file_read=_b(tb and tb in observed),
        edit_target_read=_b(edit_path and edit_path in observed),
        last_failure_type=state.get("last_failure_type") or "none",
        repeated_command=_b(tool == "Bash" and cmd and sum(1 for x in hist[-3:] if x == cmd) >= 2),
    )


def _fmt_traj(calls):
    lines = []
    for i, c in enumerate(calls[-12:], 1):
        lines.append(f"{i}. {c['tool']}: {c.get('arg','')}\n   Result: {c.get('result','')}")
    return "\n".join(lines) if lines else "none"