"""Amortize LangChain's partial-JSON parsing without buffering live UI output."""

from langchain_core.messages import AIMessageChunk
from langchain_core.messages.ai import add_ai_message_chunks

CHUNK_BATCH_SIZE = 128


class ChunkAccumulator:
    def __init__(self):
        self.response = None
        self.pending = []

    @property
    def has_output(self):
        return self.response is not None or bool(self.pending)

    @property
    def ready(self):
        return len(self.pending) >= CHUNK_BATCH_SIZE

    def append(self, chunk):
        self.pending.append(chunk)

    def flush(self):
        chunks = ([self.response] if self.response is not None else []) + self.pending
        if not chunks:
            return None
        if all(isinstance(chunk, AIMessageChunk) for chunk in chunks):
            self.response = add_ai_message_chunks(chunks[0], *chunks[1:])
        else:
            result = chunks[0]
            for chunk in chunks[1:]:
                result = result + chunk
            self.response = result
        self.pending.clear()
        return self.response