"""Attach tool results as observations on agent steps."""

from __future__ import annotations

from cc_logger.parser.classify import extract_tool_results, get_tool_result_text
from cc_logger.parser.reassemble import ClassifiedMessage
from cc_logger.types.atif import ATIFObservation, ATIFObservationResult, ATIFStep


def attach_observations(
    steps: list[ATIFStep],
    classified: list[ClassifiedMessage],
) -> None:
    """Match tool_result messages to agent steps and attach as observations.

    Modifies steps in-place.
    """
    # Build a map of tool_call_id -> step for all agent steps with tool_calls
    tool_call_to_step: dict[str, ATIFStep] = {}
    for step in steps:
        if step.tool_calls:
            for tc in step.tool_calls:
                tool_call_to_step[tc.tool_call_id] = step

    # Process all tool_result classified messages
    for cm in classified:
        if cm.type != "tool_result":
            continue
        content = cm.original.content
        if not isinstance(content, list):
            continue

        tool_results = extract_tool_results(content)
        for tr in tool_results:
            tool_use_id = tr.get("tool_use_id", "")
            step = tool_call_to_step.get(tool_use_id)
            if not step:
                continue

            is_error = tr.get("is_error", False)
            text = get_tool_result_text(tr)
            if is_error:
                text = f"ERROR: {text}"

            if step.observation is None:
                step.observation = ATIFObservation(results=[])

            step.observation.results.append(
                ATIFObservationResult(
                    source_call_id=tool_use_id,
                    content=text,
                )
            )
