"""Link subagent trajectories and export them as separate ATIF files."""

from __future__ import annotations

import json
import logging
import re
from pathlib import Path

from cc_logger.types.atif import ATIFStep, SubagentTrajectoryRef
from cc_logger.types.session import SubagentInfo

logger = logging.getLogger(__name__)


def link_subagent_refs(
    steps: list[ATIFStep],
    subagents: list[SubagentInfo],
    session_id: str,
    output_dir: Path,
    export_fn: callable,  # type: ignore[valid-type]
) -> None:
    """For Task tool_calls, find matching subagent files and attach refs.

    Also exports each subagent as a separate ATIF file.
    Modifies steps in-place.
    """
    if not subagents:
        return

    # Build agent_id -> SubagentInfo map
    agent_map: dict[str, SubagentInfo] = {sa.agent_id: sa for sa in subagents}

    for step in steps:
        if not step.tool_calls or not step.observation:
            continue

        for i, tc in enumerate(step.tool_calls):
            if tc.function_name != "Task":
                continue

            # Find the agent ID from the observation result
            if i >= len(step.observation.results):
                continue

            obs_result = step.observation.results[i]
            content = obs_result.content or ""

            # Extract agent ID from the observation content
            # Pattern: "agentId: abc1234"
            match = re.search(r"agentId:\s*([a-f0-9]+)", content)
            if not match:
                continue

            agent_id = match.group(1)
            subagent = agent_map.get(agent_id)
            if not subagent:
                logger.warning("Subagent %s not found for Task tool_call", agent_id)
                continue

            # Export the subagent trajectory
            subagent_output_dir = output_dir / session_id / "subagents"
            subagent_output_dir.mkdir(parents=True, exist_ok=True)
            subagent_path = subagent_output_dir / f"agent-{agent_id}.json"

            try:
                trajectory = export_fn(subagent.jsonl_path, subagent_session_id=f"{session_id}/agent-{agent_id}")
                subagent_path.write_text(json.dumps(trajectory, indent=2))

                # Attach ref
                rel_path = f"{session_id}/subagents/agent-{agent_id}.json"
                obs_result.subagent_trajectory_ref = [
                    SubagentTrajectoryRef(
                        session_id=f"{session_id}/agent-{agent_id}",
                        trajectory_path=rel_path,
                    )
                ]
            except Exception as e:
                logger.warning("Failed to export subagent %s: %s", agent_id, e)
