"""Reassemble streamed assistant responses from multiple JSONL lines."""

from __future__ import annotations

from dataclasses import dataclass, field

from cc_logger.types.cc import (
    CCMessage,
    CCUsage,
    ContentBlock,
    TextBlock,
    ThinkingBlock,
    ToolUseBlock,
)


@dataclass
class ReassembledResponse:
    """A single assistant API response, reassembled from multiple JSONL lines."""

    message_id: str
    model: str | None = None
    content_blocks: list[ContentBlock] = field(default_factory=list)
    usage: CCUsage = field(default_factory=CCUsage)
    timestamp: str | None = None
    uuid: str | None = None  # UUID of the first line in the group

    @property
    def thinking_blocks(self) -> list[ThinkingBlock]:
        return [b for b in self.content_blocks if isinstance(b, ThinkingBlock)]

    @property
    def text_blocks(self) -> list[TextBlock]:
        return [b for b in self.content_blocks if isinstance(b, TextBlock)]

    @property
    def tool_use_blocks(self) -> list[ToolUseBlock]:
        return [b for b in self.content_blocks if isinstance(b, ToolUseBlock)]

    @property
    def reasoning_text(self) -> str | None:
        texts = [b.thinking for b in self.thinking_blocks if b.thinking]
        return "\n\n".join(texts) if texts else None

    @property
    def message_text(self) -> str | None:
        texts = [b.text for b in self.text_blocks if b.text]
        return "\n\n".join(texts) if texts else None


@dataclass
class ClassifiedMessage:
    """A message in the conversation with its classification."""

    type: str  # "user", "agent", "system", "tool_result", "compact_summary", "skip"
    original: CCMessage
    reassembled: ReassembledResponse | None = None  # only for "agent"


def reassemble_assistant_groups(messages: list[CCMessage]) -> list[ClassifiedMessage]:
    """Process a linear sequence of CCMessages, grouping assistant chunks and classifying all messages.

    Returns a list of ClassifiedMessage objects ready for ATIF mapping.
    """
    result: list[ClassifiedMessage] = []
    i = 0

    while i < len(messages):
        msg = messages[i]

        if msg.type == "assistant":
            # Check if synthetic
            if msg.model == "<synthetic>":
                response = ReassembledResponse(
                    message_id=msg.message_id or "",
                    model=msg.model,
                    content_blocks=list(msg.content_blocks),
                    usage=msg.usage,
                    timestamp=msg.timestamp,
                    uuid=msg.uuid,
                )
                result.append(ClassifiedMessage(type="agent", original=msg, reassembled=response))
                i += 1
                continue

            # Group consecutive assistant lines with same message.id
            group_id = msg.message_id
            group: list[CCMessage] = [msg]
            j = i + 1
            while j < len(messages):
                next_msg = messages[j]
                if (
                    next_msg.type == "assistant"
                    and next_msg.message_id == group_id
                    and next_msg.model != "<synthetic>"
                    and group_id is not None
                ):
                    group.append(next_msg)
                    j += 1
                else:
                    break

            # Build reassembled response from group
            all_blocks: list[ContentBlock] = []
            for g in group:
                all_blocks.extend(g.content_blocks)

            response = ReassembledResponse(
                message_id=group_id or "",
                model=group[0].model,
                content_blocks=all_blocks,
                usage=group[-1].usage,  # last chunk has most complete usage
                timestamp=group[0].timestamp,
                uuid=group[0].uuid,
            )
            result.append(ClassifiedMessage(type="agent", original=group[0], reassembled=response))
            i = j

        elif msg.type == "user":
            if msg.is_compact_summary:
                result.append(ClassifiedMessage(type="compact_summary", original=msg))
            elif isinstance(msg.content, list):
                # Tool result message
                result.append(ClassifiedMessage(type="tool_result", original=msg))
            else:
                # External user message
                result.append(ClassifiedMessage(type="user", original=msg))
            i += 1

        elif msg.type == "system":
            subtype = msg.subtype
            if subtype in ("compact_boundary", "microcompact_boundary", "api_error", "local_command"):
                result.append(ClassifiedMessage(type="system", original=msg))
            else:
                result.append(ClassifiedMessage(type="skip", original=msg))
            i += 1

        else:
            # progress, file-history-snapshot, queue-operation → skip
            result.append(ClassifiedMessage(type="skip", original=msg))
            i += 1

    return result
