from typing import Any

from zrb.context.any_context import zrb_print
from zrb.llm.config.limiter import LLMLimiter
from zrb.llm.config.limiter import llm_limiter as default_llm_limiter
from zrb.llm.summarizer.message_converter import message_to_text
from zrb.llm.summarizer.text_summarizer import summarize_text_plain
from zrb.util.cli.style import stylize_error, stylize_yellow


async def chunk_and_summarize(
    messages: list[Any],
    agent: Any,
    limiter: "LLMLimiter",
    token_threshold: int,
    include_last_user_intent_instruction: bool = False,
) -> str:
    """Break history into chunks and summarize them iteratively."""
    # Pre-calculate texts
    history_texts = []
    for m in messages:
        try:
            history_texts.append(message_to_text(m))
        except Exception as e:
            zrb_print(
                stylize_error(f"  Error converting message to text: {e}"), plain=True
            )
            history_texts.append(str(m))
    summaries = []
    current_chunk = []
    current_chunk_tokens = 0
    chunk_token_limit = max(1, int(token_threshold * 0.9))

    async def _flush_chunk(is_last_chunk: bool = False):
        if not current_chunk:
            return
        zrb_print(
            stylize_yellow(f"  Compressing chunk of {len(current_chunk)} messages..."),
            plain=True,
        )
        try:
            # Use summarize_text_plain for individual chunks to handle their potential size safely
            chunk_content = "\n".join(current_chunk)

            if is_last_chunk and include_last_user_intent_instruction:
                chunk_content += (
                    "\n\nIMPORTANT: The last part of this history contains the user's latest request. "
                    "Ensure the <state_snapshot> explicitly captures this request in <overall_goal> or <task_state> "
                    "so the next agent knows exactly what to do."
                )

            summary = await summarize_text_plain(
                chunk_content, agent, limiter, token_threshold
            )
            summaries.append(summary)
        except Exception as e:
            zrb_print(stylize_error(f"  Error summarizing chunk: {e}"), plain=True)
            raise e

    for text in history_texts:
        text_tokens = limiter.count_tokens(text)
        if current_chunk and (current_chunk_tokens + text_tokens > chunk_token_limit):
            await _flush_chunk(is_last_chunk=False)
            current_chunk = []
            current_chunk_tokens = 0
        current_chunk.append(text)
        current_chunk_tokens += text_tokens
    await _flush_chunk(is_last_chunk=True)

    return "\n\n".join(summaries) if summaries else "[No summaries generated]"


async def consolidate_summaries(
    summary_text: str,
    agent: Any,
    conversational_token_threshold: int,
    has_multiple_snapshots: bool,
) -> str:
    if has_multiple_snapshots:
        zrb_print(stylize_yellow("  Consolidating multiple snapshots..."), plain=True)
    else:
        zrb_print(
            stylize_yellow("  Aggressively re-compressing summary..."),
            plain=True,
        )

    # Use summarize_text_plain to safely consolidate even a large block of summaries
    prompt = (
        "Consolidate the following conversation state snapshots into a single, cohesive "
        "<state_snapshot>.\n"
        "IMPORTANT: Be extremely concise and dense. Your goal is to fit all critical "
        f"knowledge into less than {conversational_token_threshold // 2} tokens.\n"
        f"{summary_text}"
    )
    return await summarize_text_plain(
        prompt, agent, default_llm_limiter, conversational_token_threshold
    )
