"""Assemble a complete ATIF trajectory from a CC session JSONL."""

from __future__ import annotations

from pathlib import Path
from typing import Any

from cc_logger.mapper.metrics import compute_final_metrics
from cc_logger.mapper.observations import attach_observations
from cc_logger.mapper.step_builder import build_steps
from cc_logger.mapper.subagents import link_subagent_refs
from cc_logger.parser.jsonl_reader import read_jsonl
from cc_logger.parser.reassemble import reassemble_assistant_groups
from cc_logger.parser.tree_builder import linearize_conversation
from cc_logger.types.atif import ATIFAgent, ATIFTrajectory
from cc_logger.types.cc import parse_message
from cc_logger.types.session import SubagentInfo


def export_session(
    jsonl_path: Path,
    subagents: list[SubagentInfo] | None = None,
    output_dir: Path | None = None,
    subagent_session_id: str | None = None,
) -> dict[str, Any]:
    """Export a CC session JSONL file to an ATIF trajectory dict.

    Args:
        jsonl_path: Path to the session JSONL file
        subagents: List of subagent info for this session
        output_dir: Output directory (needed for subagent file export)
        subagent_session_id: Override session ID (for subagent exports)

    Returns:
        ATIF trajectory as a dict ready for JSON serialization
    """
    # 1. Read all JSONL lines
    raw_messages = list(read_jsonl(jsonl_path))

    # 2. Linearize conversation tree
    linear = linearize_conversation(raw_messages)

    # 3. Parse into typed messages
    parsed = [parse_message(m) for m in linear]

    # 4. Reassemble assistant responses and classify all messages
    classified = reassemble_assistant_groups(parsed)

    # 5. Build ATIF steps
    steps = build_steps(classified)

    # 6. Attach tool result observations
    attach_observations(steps, classified)

    # 7. Link subagent trajectories
    if subagents and output_dir:
        session_id = subagent_session_id or jsonl_path.stem
        link_subagent_refs(
            steps=steps,
            subagents=subagents,
            session_id=session_id,
            output_dir=output_dir,
            export_fn=lambda path, **kw: export_session(path, subagent_session_id=kw.get("subagent_session_id")),
        )

    # 8. Extract agent config
    agent = _extract_agent_config(parsed)

    # 9. Compute final metrics
    final_metrics = compute_final_metrics(steps)

    # 10. Build trajectory
    session_id = subagent_session_id or jsonl_path.stem
    trajectory = ATIFTrajectory(
        session_id=session_id,
        agent=agent,
        steps=steps,
        final_metrics=final_metrics,
    )

    return trajectory.model_dump(exclude_none=True)


def _extract_agent_config(messages: list) -> ATIFAgent:
    """Extract agent name, version, and model from messages."""
    version = "unknown"
    model_name = None

    for msg in messages:
        if msg.version and version == "unknown":
            version = msg.version
        if msg.type == "assistant" and msg.model and msg.model != "<synthetic>":
            if model_name is None:
                model_name = msg.model
        if version != "unknown" and model_name:
            break

    return ATIFAgent(
        name="claude-code",
        version=version,
        model_name=model_name,
    )
